changes from apache#8131 (#228)

neo-ai · Sep 28, 2021 · 561ac46 · 561ac46
1 parent 5203272
commit 561ac46
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git 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
@@ -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;