From 51253ac14aeed52ffa6402dbc59b98d6ea72e85c Mon Sep 17 00:00:00 2001 From: Dun Liang Date: Sun, 7 Jun 2020 12:24:46 +0800 Subject: [PATCH] fix model test without pth --- python/jittor/test/test_models.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/python/jittor/test/test_models.py b/python/jittor/test/test_models.py index cc42dfc3..55d61aa8 100644 --- 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()