fix test_raise

apache · Feb 12, 2021 · 2bf84de · 2bf84de
1 parent 7ede696
commit 2bf84de
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/python-pytest/onnx/test_operators.py b/tests/python-pytest/onnx/test_operators.py
@@ -1139,6 +1139,6 @@ def test_onnx_export_take(tmp_path, dtype, axis, mode):
 @pytest.mark.parametrize('axis', [-3, -2, -1, 0, 1, 2])
 def test_onnx_export_take_raise(tmp_path, dtype, axis):
     x = mx.nd.random.normal(0, 10, (3, 4, 5)).astype(dtype)
-    y = mx.random.randint(0, 4, (6, 7)).astype(dtype)
+    y = mx.random.randint(0, 3, (6, 7)).astype(dtype)
     M = def_model('take', axis=axis, mode='raise')
     op_export_test('take', M, [x, y], tmp_path)