passing mat size as arg

apache · Jul 23, 2020 · 7dc2004 · 7dc2004
1 parent 71a8778
commit 7dc2004
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/python/mxnet/test_utils.py b/python/mxnet/test_utils.py
@@ -298,9 +298,9 @@ def create_vector(size, dtype=np.int64):
     return a
 
 # For testing Large Square Matrix with total size > 2^32 elements
-def get_large_identity_mat():
-    A = nd.zeros((LARGE_SQ_X, LARGE_SQ_X))
-    for i in range(LARGE_SQ_X):
+def get_large_identity_mat(size):
+    A = nd.zeros((size, size))
+    for i in range(size):
         A[i,i] = 1
     return A
 

diff --git a/tests/nightly/test_large_array.py b/tests/nightly/test_large_array.py
@@ -1204,7 +1204,7 @@ def run_trsm(inp):
             out = trsm(inp, inp)
         return inp.grad, out
 
-    A = get_large_identity_mat()
+    A = get_large_identity_mat(LARGE_SQ_X)
 
     grad, out = run_det(A)
     check_diag(grad, 0)