Add Handling of Zero Len Arguments (apache#6923)

* Update tensorrt.py * Update tensorrt.py * Update tensorrt.py
neo-ai · Dec 4, 2020 · d0d541a · d0d541a
1 parent bebeb52
commit d0d541a
Showing 1 changed file with 4 additions and 0 deletions.
diff --git 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