From 67bf1a0702c68f975b8711cc2b62000d0e37d800 Mon Sep 17 00:00:00 2001 From: Ritwik Das Date: Wed, 18 Nov 2020 09:46:26 -0800 Subject: [PATCH] Add Handling of Zero Len Arguments (#6923) * Update tensorrt.py * Update tensorrt.py * Update tensorrt.py --- 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 739d49c412e8..44336073d842 100644 --- a/python/tvm/relay/op/contrib/tensorrt.py +++ b/python/tvm/relay/op/contrib/tensorrt.py @@ -292,6 +292,10 @@ def add_annotate_fn(expr): # pylint: disable=unused-variable """Check if add is supported by TensorRT.""" args = expr.args + # RelayVM + TRT doesn't support scalar addition yet. + for arg in args: + if not arg.checked_type.shape: + return False if any([x.checked_type.dtype != "float32" for x in args]): logger.info("Only float32 inputs are supported for TensorRT.") return False