fix: Fix for torch scripted module faiure with DLFW

Signed-off-by: Anurag Dixit <anuragd@nvidia.com>
pytorch · Nov 23, 2021 · 88c02d9 · 88c02d9
1 parent a10613e
commit 88c02d9
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/tests/py/test_api.py b/tests/py/test_api.py
@@ -30,12 +30,13 @@ def test_compile_traced(self):
         self.assertTrue(same < 2e-2)
 
     def test_compile_script(self):
-        trt_mod = torchtrt.ts.compile(self.scripted_model,
+        with torch.no_grad():
+            trt_mod = torchtrt.ts.compile(self.scripted_model,
                                       inputs=[self.input],
                                       device=torchtrt.Device(gpu_id=0),
                                       enabled_precisions={torch.float})
-        same = (trt_mod(self.input) - self.scripted_model(self.input)).abs().max()
-        self.assertTrue(same < 2e-2)
+            same = (trt_mod(self.input) - self.scripted_model(self.input)).abs().max()
+            self.assertTrue(same < 2e-2)
 
     def test_compile_global(self):
         trt_mod = torchtrt.compile(self.scripted_model,
@@ -46,12 +47,13 @@ def test_compile_global(self):
         self.assertTrue(same < 2e-2)
 
     def test_compile_global_nn_mod(self):
-        trt_mod = torchtrt.compile(self.model,
+        with torch.no_grad():
+            trt_mod = torchtrt.compile(self.model,
                                    inputs=[self.input],
                                    device=torchtrt.Device(gpu_id=0),
                                    enabled_precisions={torch.float})
-        same = (trt_mod(self.input) - self.scripted_model(self.input)).abs().max()
-        self.assertTrue(same < 2e-2)
+            same = (trt_mod(self.input) - self.scripted_model(self.input)).abs().max()
+            self.assertTrue(same < 2e-2)
 
     def test_from_torch_tensor(self):
         compile_spec = {