fix default axis behavior

AndrewZhaoLuo · Sep 23, 2021 · a718b2a · a718b2a
1 parent 8ef89ea
commit a718b2a
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py
@@ -1461,7 +1461,8 @@ def run_calculation(cls, tensor, axes):
 
     @classmethod
     def _impl_v1(cls, inputs, attr, params):
-        axes = sorted(attr["axes"])
+        shape = infer_shape(inputs[0])
+        axes = sorted(attr.get("axes", [i for i, size in enumerate(shape) if size == 1]))
         return cls.run_calculation(inputs[0], axes)
 
     @classmethod