fix test.

apache · May 29, 2019 · c283601 · c283601
1 parent d1403d7
commit c283601
Showing 1 changed file with 6 additions and 7 deletions.
diff --git a/tests/python/unittest/test_gluon.py b/tests/python/unittest/test_gluon.py
@@ -2735,13 +2735,12 @@ def __init__(self, **kwargs):
         def forward(self, x):
             return self.dense(x)
 
-    mx.set_np_shape(True)
-    z = mx.nd.zeros((2,2016))
-    print(z.shape)
-    foo = Foo()
-    foo.initialize()
-    print(foo(z).shape)
-    mx.set_np_shape(False)
+    with mx.np_shape(True):
+        z = mx.nd.zeros((2,2016))
+        print(z.shape)
+        foo = Foo()
+        foo.initialize()
+        print(foo(z).shape)
 
 if __name__ == '__main__':
     import nose