wrap forward passes with torch.no_grad() (#19438)

huggingface · Oct 10, 2022 · 870a954 · 870a954
1 parent 692c5be
commit 870a954
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/models/roformer/test_modeling_roformer.py b/tests/models/roformer/test_modeling_roformer.py
@@ -457,7 +457,8 @@ class RoFormerModelIntegrationTest(unittest.TestCase):
     def test_inference_masked_lm(self):
         model = RoFormerForMaskedLM.from_pretrained("junnyu/roformer_chinese_base")
         input_ids = torch.tensor([[0, 1, 2, 3, 4, 5]])
-        output = model(input_ids)[0]
+        with torch.no_grad():
+            output = model(input_ids)[0]
 
         # TODO Replace vocab size
         vocab_size = 50000