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

kashif · Oct 21, 2022 · 0ffbc82 · 0ffbc82
1 parent 5b1d462
commit 0ffbc82
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/models/flaubert/test_modeling_flaubert.py b/tests/models/flaubert/test_modeling_flaubert.py
@@ -437,7 +437,8 @@ class FlaubertModelIntegrationTest(unittest.TestCase):
     def test_inference_no_head_absolute_embedding(self):
         model = FlaubertModel.from_pretrained("flaubert/flaubert_base_cased")
         input_ids = torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]])
-        output = model(input_ids)[0]
+        with torch.no_grad():
+            output = model(input_ids)[0]
         expected_shape = torch.Size((1, 11, 768))
         self.assertEqual(output.shape, expected_shape)
         expected_slice = torch.tensor(