feat(): support for adaptive_avg_pool2d plugin

Signed-off-by: Abhiram Iyer <[email protected]> Signed-off-by: Abhiram Iyer <[email protected]>
pytorch · Jun 18, 2020 · 52be580 · 52be580
1 parent fa227b0
commit 52be580
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/core/conversion/converters/impl/plugins/interpolate_plugin.cpp b/core/conversion/converters/impl/plugins/interpolate_plugin.cpp
@@ -175,6 +175,8 @@ int InterpolatePlugin::enqueue(const nvinfer1::PluginTensorDesc* inputDesc, cons
         at::upsample_bilinear2d_out(output, input, {size[0], size[1]}, align_corners);
     } else if (mode == "trilinear") {
         at::upsample_trilinear3d_out(output, input, {size[0], size[1], size[2]}, align_corners);
+    } else if (mode == "adaptive_pool2d") {
+        at::adaptive_avg_pool2d_out(output, input, {size[0], size[1]});
     }
 
     cudaEvent_t torch_event;