diff --git a/python/tvm/relay/op/contrib/tensorrt.py b/python/tvm/relay/op/contrib/tensorrt.py index a4148e4fe4f2..dfdda633b4fa 100644 --- a/python/tvm/relay/op/contrib/tensorrt.py +++ b/python/tvm/relay/op/contrib/tensorrt.py @@ -733,8 +733,12 @@ def pad_annotate_fn(expr): # pylint: disable=unused-variable if attrs.pad_mode != "constant": logger.info("nn.pad: pad mode is %s but must be constant.", attrs.pad_mode) return False - if float(attrs.pad_value) != 0.0: - logger.info("nn.pad: pad value is %f but must be 0.0.", float(attrs.pad_value)) + if ( + not isinstance(args[1], relay.Constant) + or len(args[1].checked_type.shape) != 0 + or args[1].data.numpy().item() != 0.0 + ): + logger.info("nn.pad: pad value is %s but must be 0.0.", args[1]) return False if len(attrs.pad_width) not in [4, 5]: logger.info("nn.pad: can only pad 4D or 5D inputs") diff --git a/src/runtime/contrib/tensorrt/tensorrt_ops.cc b/src/runtime/contrib/tensorrt/tensorrt_ops.cc index 7197172d73db..020ef866692a 100644 --- a/src/runtime/contrib/tensorrt/tensorrt_ops.cc +++ b/src/runtime/contrib/tensorrt/tensorrt_ops.cc @@ -1057,7 +1057,7 @@ class ReshapeOpConverter : public TensorRTOpConverter { class PadOpConverter : public TensorRTOpConverter { public: - PadOpConverter() : TensorRTOpConverter({kTensor}) {} + PadOpConverter() : TensorRTOpConverter({kTensor, kWeight}) {} void Convert(TensorRTOpConverterParams* params) const { auto input = params->inputs.at(0).tensor;