From a9ccf25edc815bc30498c2235ae52ca8ae214d1b Mon Sep 17 00:00:00 2001 From: Trevor Morris Date: Tue, 18 May 2021 23:11:33 -0700 Subject: [PATCH] Only allow 4d or 5d inputs to TRT nn.pad (#8073) --- python/tvm/relay/op/contrib/tensorrt.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/python/tvm/relay/op/contrib/tensorrt.py b/python/tvm/relay/op/contrib/tensorrt.py index 266c4987b27a..cbe6a22f4a4d 100644 --- a/python/tvm/relay/op/contrib/tensorrt.py +++ b/python/tvm/relay/op/contrib/tensorrt.py @@ -726,11 +726,15 @@ def pad_annotate_fn(expr): # pylint: disable=unused-variable if float(attrs.pad_value) != 0.0: logger.info("nn.pad: pad value is %f but must be 0.0.", float(attrs.pad_value)) return False + if len(attrs.pad_width) not in [4, 5]: + logger.info("nn.pad: can only pad 4D or 5D inputs") + return False if any([x != 0 for x in attrs.pad_width[0]]) or any([x != 0 for x in attrs.pad_width[1]]): logger.info("nn.pad: can't pad batch or channel dimensions.") return False if len(attrs.pad_width) == 5 and any([x != 0 for x in attrs.pad_width[2]]): logger.info("nn.pad: can only pad last two dimensions for 5D inputs.") + return False return True