nit

Author: kazhang

test/test_models.py

@@ -133,12 +133,10 @@ def get_export_import_copy(m):
             return imported
 
         m_import = get_export_import_copy(m)
-        with torch.no_grad():
-            with freeze_rng_state():
-                results = m(*args)
-        with torch.no_grad():
-            with freeze_rng_state():
-                results_from_imported = m_import(*args)
+        with torch.no_grad(), freeze_rng_state():
+            results = m(*args)
+        with torch.no_grad(), freeze_rng_state():
+            results_from_imported = m_import(*args)
         tol = 3e-4
         torch.testing.assert_close(results, results_from_imported, atol=tol, rtol=tol)
 
@@ -158,12 +156,10 @@ def get_export_import_copy(m):
 
     sm = torch.jit.script(nn_module)
 
-    with torch.no_grad():
-        with freeze_rng_state():
+    with torch.no_grad(), freeze_rng_state():
             eager_out = nn_module(*args)
 
-    with torch.no_grad():
-        with freeze_rng_state():
+    with torch.no_grad(), freeze_rng_state():
             script_out = sm(*args)
             if unwrapper:
                 script_out = unwrapper(script_out)