From 7794c78c445b78f802dd076904eeff907f9a8c5d Mon Sep 17 00:00:00 2001 From: Abhiram Iyer Date: Thu, 18 Jun 2020 18:23:33 -0700 Subject: [PATCH] fix(): fixed interpolate_plugin to handle dynamically sized inputs for adaptive_pool2d Signed-off-by: Abhiram Iyer Signed-off-by: Abhiram Iyer --- .../converters/impl/plugins/interpolate_plugin.cpp | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/core/conversion/converters/impl/plugins/interpolate_plugin.cpp b/core/conversion/converters/impl/plugins/interpolate_plugin.cpp index 2da05a70e5..7ebf5d18ea 100755 --- a/core/conversion/converters/impl/plugins/interpolate_plugin.cpp +++ b/core/conversion/converters/impl/plugins/interpolate_plugin.cpp @@ -156,8 +156,16 @@ size_t InterpolatePlugin::getWorkspaceSize(const nvinfer1::PluginTensorDesc* inp int InterpolatePlugin::enqueue(const nvinfer1::PluginTensorDesc* inputDesc, const nvinfer1::PluginTensorDesc* outputDesc, const void *const *inputs, void *const *outputs, void *workspace, cudaStream_t stream) { + at::Tensor input; + + if (mode == "adaptive_pool2d") { + // use dynamically inferred input shape (for pooling) + input = at::from_blob((void*) inputs[0], util::toVec(inputDesc->dims), [](void*){}, tensor_options); + } else { + // use precomputed input shape (for interpolation/upsampling) + input = at::from_blob((void*) inputs[0], in_shape, [](void*){}, tensor_options); + } - at::Tensor input = at::from_blob((void*) inputs[0], in_shape, [](void*){}, tensor_options); at::Tensor output = at::from_blob(outputs[0], out_shape, [](void*){}, tensor_options); at::cuda::CUDAStream torch_stream = at::cuda::getStreamFromPool();