Fixed onnx test failures when run on a cpu backend

apache · Aug 15, 2019 · b870cc9 · b870cc9
1 parent 7eb1f35
commit b870cc9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py
@@ -1082,7 +1082,7 @@ def check_torch_conversion(model, input_size):
     # Set verbose=True for more output
     torch.onnx.export(model(), dummy_input, file_name, export_params=True, verbose=False)
     onnx_model = onnx.load(file_name)
-    shapes = { '0' : input_size }
+    shapes = { onnx_model.graph.input[0].name : input_size }
     expr, params = relay.frontend.from_onnx(onnx_model, shape=shapes)
 
 def test_resnet():