fix model test without pth

2020-06-07 12:24:46 +08:00 · 2020-06-07 12:24:46 +08:00 · 51253ac14a
parent 8ffad0230d
commit 51253ac14a
1 changed files with 4 additions and 1 deletions
--- a/python/jittor/test/test_models.py
+++ b/python/jittor/test/test_models.py
@ -61,8 +61,11 @@ class test_models(unittest.TestCase):

    @unittest.skipIf(not jt.has_cuda, "Cuda not found")
    @jt.flag_scope(use_cuda=1)
-    @torch.no_grad()
    def test_models(self):
+        with torch.no_grad():
+            self.run_models()
+
+    def run_models(self):
        def to_cuda(x):
            if jt.has_cuda:
                return x.cuda()