gather_nd: update test

tests/python/unittest/test_operator.py

@@ -7167,15 +7167,26 @@ def check(data, idx):
             assert (mx.nd._internal._backward_gather_nd(data, idx, shape=(1,)).asscalar() == data.asnumpy().sum())
 
 @with_seed()
-def test_gather_nd_bound():
+def test_gather_nd_check_bound():
     # check if indices is out of bound
     data = mx.nd.array([[0, 1, 2], [3, 4, 5]])
     indices1 = mx.nd.array([[0, 1, 0], [0, 1, 3]])
     indices2 = mx.nd.array([[0, 1, 0], [0, 1, -5]])
-    assertRaises(IndexError, mx.nd.gather_nd, data, indices1)
-    # IndexError: Check failed: false: index 3 is out of bounds for axis 1 with size 3
-    assertRaises(IndexError, mx.nd.gather_nd, data, indices2)
-    # IndexError: index -5 is out of bounds for axis 1 with size 3
+    try:
+        mx.nd.gather_nd(data, indices1)
+        mx.nd.waitall()
+    except IndexError:
+            # skip errors since the test is supposed to raise error
+            # IndexError: index 3 is out of bounds for axis 1 with size 3
+            pass
+
+    try:
+        mx.nd.gather_nd(data, indices2)
+        mx.nd.waitall()
+    except IndexError:
+            # skip errors since the test is supposed to raise error
+            # IndexError: index -5 is out of bounds for axis 1 with size 3
+            pass
 
     # check if the negative indices are wrapped correctly
     indices1 = mx.nd.array([[0, 1, -1], [0, 1, -2]])