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

870a9542 · Partho · GitHub · 692c5be7 · 870a9542
Unverified Commit 870a9542 authored Oct 11, 2022 by Partho Committed by GitHub Oct 10, 2022
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

tests/models/roformer/test_modeling_roformer.py tests/models/roformer/test_modeling_roformer.py +2 -1

No files found.
--- a/tests/models/roformer/test_modeling_roformer.py
+++ b/tests/models/roformer/test_modeling_roformer.py
@@ -457,6 +457,7 @@ 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]])
+        with torch.no_grad():
            output = model(input_ids)[0]

        # TODO Replace vocab size