adding asnumpy() to output of gather(implicitly called) to fix gather…

… test in large vector and tensor tests (#17290)
apache · Jan 14, 2020 · 71e70f2 · 71e70f2
1 parent 9f2e73f
commit 71e70f2
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/tests/nightly/test_large_array.py b/tests/nightly/test_large_array.py
@@ -1672,10 +1672,10 @@ def test_gather():
     idx = mx.nd.random.randint(0, LARGE_X, SMALL_X)
     # Calls gather_nd internally
     tmp = arr[idx]
-    assert np.sum(tmp[0] == 1) == SMALL_Y
+    assert np.sum(tmp[0].asnumpy() == 1) == SMALL_Y
     # Calls gather_nd internally
     arr[idx] += 1
-    assert np.sum(arr[idx[0]] == 2) == SMALL_Y
+    assert np.sum(arr[idx[0]].asnumpy() == 2) == SMALL_Y
 
 
 def test_binary_broadcast():

diff --git a/tests/nightly/test_large_vector.py b/tests/nightly/test_large_vector.py
@@ -1049,10 +1049,10 @@ def test_gather():
     idx = mx.nd.random.randint(0, LARGE_X, 10, dtype=np.int64)
     # Calls gather_nd internally
     tmp = arr[idx]
-    assert np.sum(tmp == 1) == 10
+    assert np.sum(tmp.asnumpy() == 1) == 10
     # Calls gather_nd internally
     arr[idx] += 1
-    assert np.sum(arr[idx] == 2) == 10
+    assert np.sum(arr[idx].asnumpy() == 2) == 10
 
 
 def test_infer_shape():