From 70836d06616dd2eb76e3b19328ddd33f7f19572b Mon Sep 17 00:00:00 2001 From: Dong Xu Date: Tue, 27 Feb 2024 15:09:53 +0800 Subject: [PATCH] [TRT] Add fallback to hand-written ROIAlign for TRT prev to 8.4 --- .../network/tensorrt/layer_builder/roialign_layer_builder.cc | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/source/tnn/network/tensorrt/layer_builder/roialign_layer_builder.cc b/source/tnn/network/tensorrt/layer_builder/roialign_layer_builder.cc index 9505116ff..06e9a8892 100644 --- a/source/tnn/network/tensorrt/layer_builder/roialign_layer_builder.cc +++ b/source/tnn/network/tensorrt/layer_builder/roialign_layer_builder.cc @@ -45,6 +45,7 @@ nvinfer1::DataType RoiAlignTRTPluginLayerBuilder::getOutputDataType(int index, c } ILayer* RoiAlignTRTPluginLayerBuilder::AddToNetwork(INetworkDefinition* network) noexcept { +#if NV_TENSORRT_MAJOR * 10 + NV_TENSORRT_MINOR >= 84 auto layer_param = dynamic_cast(param_); if (!layer_param) { LOGE("RoiAlignTRTPluginLayerBuilder: Unable to get layer param."); @@ -86,6 +87,9 @@ ILayer* RoiAlignTRTPluginLayerBuilder::AddToNetwork(INetworkDefinition* network) } return layer; +#else + return TensorRTPluginLayerBuilder::AddToNetwork(network); +#endif } DimsExprs RoiAlignTRTPluginLayerBuilder::getOutputDimensions(int index, const nvinfer1::DimsExprs* inputs,