From 0ffbc822493b58cbdda4d9eddc6ad047b9d425f3 Mon Sep 17 00:00:00 2001 From: Partho Date: Tue, 11 Oct 2022 00:34:10 +0530 Subject: [PATCH] wrap forward passes with torch.no_grad() (#19412) --- tests/models/flaubert/test_modeling_flaubert.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/models/flaubert/test_modeling_flaubert.py b/tests/models/flaubert/test_modeling_flaubert.py index da29cac6dd588c..b4150a65823575 100644 --- 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(