wrap forward passes with torch.no_grad() (huggingface#19413)

amyeroberts · Oct 18, 2022 · 5b02525 · 5b02525
1 parent d3c0436
commit 5b02525
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/tests/models/fnet/test_modeling_fnet.py b/tests/models/fnet/test_modeling_fnet.py
@@ -493,7 +493,8 @@ def test_inference_for_masked_lm(self):
         model.to(torch_device)
 
         input_ids = torch.tensor([[0, 1, 2, 3, 4, 5]], device=torch_device)
-        output = model(input_ids)[0]
+        with torch.no_grad():
+            output = model(input_ids)[0]
 
         vocab_size = 32000
 
@@ -536,7 +537,8 @@ def test_inference_for_next_sentence_prediction(self):
         model.to(torch_device)
 
         input_ids = torch.tensor([[0, 1, 2, 3, 4, 5]], device=torch_device)
-        output = model(input_ids)[0]
+        with torch.no_grad():
+            output = model(input_ids)[0]
 
         expected_shape = torch.Size((1, 2))
         self.assertEqual(output.shape, expected_shape)
@@ -551,7 +553,8 @@ def test_inference_model(self):
         model.to(torch_device)
 
         input_ids = torch.tensor([[0, 1, 2, 3, 4, 5]], device=torch_device)
-        output = model(input_ids)[0]
+        with torch.no_grad():
+            output = model(input_ids)[0]
 
         expected_shape = torch.Size((1, 6, model.config.hidden_size))
         self.assertEqual(output.shape, expected_shape)