fix test shape

apache · Dec 18, 2020 · 67321f0 · 67321f0
1 parent 2d8dc74
commit 67321f0
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/python-pytest/onnx/test_operators.py b/tests/python-pytest/onnx/test_operators.py
@@ -261,8 +261,8 @@ def test_onnx_export_softmax(tmp_path, dtype):
     M1 = def_model('softmax')
     op_export_test('softmax_1', M1, [x], tmp_path)
     M2 = def_model('softmax', use_length=True, axis=0)
-    l2 = mx.nd.array([[2,0,3,1],[1,3,2,0], [0,0,0,1]], dtype=int)
+    l2 = mx.nd.array([[2,0,2,1],[1,1,2,1], [0,0,0,1]], dtype=int)
     op_export_test('softmax_2', M2, [x, l2], tmp_path)
     M3 = def_model('softmax', use_length=True, axis=-1, temperature=0.5)
-    l3 = mx.nd.array([[2,0,1],[0,0,0]], dtype=int)
+    l3 = mx.nd.array([[2,0,4],[0,0,0]], dtype=int)
     op_export_test('softmax_3', M3, [x, l3], tmp_path)