Raise error if key cannot be found in graph runtime inputs

apache · Jul 27, 2020 · 321cae2 · 321cae2
1 parent 5828501
commit 321cae2
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/python/tvm/contrib/graph_runtime.py b/python/tvm/contrib/graph_runtime.py
@@ -152,7 +152,10 @@ def set_input(self, key=None, value=None, **params):
            Additional arguments
         """
         if key is not None:
-            self._get_input(key).copyfrom(value)
+            v = self._get_input(key)
+            if v is None:
+                raise RuntimeError("Could not find '%s' in graph's inputs" % key)
+            v.copyfrom(value)
 
         if params:
             # upload big arrays first to avoid memory issue in rpc mode