From cc1d0685b313132eec0c188fd544925ad45ef4b6 Mon Sep 17 00:00:00 2001 From: Katie Le <54815905+katiele47@users.noreply.github.com> Date: Wed, 8 Feb 2023 08:00:52 -0500 Subject: [PATCH] Wrap RemBert integration test forward passes with torch.no_grad() (#21503) added with torch.no_grad() to the integration tests and applied make style Co-authored-by: Bibi --- tests/models/rembert/test_modeling_rembert.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/models/rembert/test_modeling_rembert.py b/tests/models/rembert/test_modeling_rembert.py index b431aff86e..180b71b17a 100644 --- a/tests/models/rembert/test_modeling_rembert.py +++ b/tests/models/rembert/test_modeling_rembert.py @@ -464,7 +464,8 @@ class RemBertModelIntegrationTest(unittest.TestCase): model = RemBertModel.from_pretrained("google/rembert") input_ids = torch.tensor([[312, 56498, 313, 2125, 313]]) segment_ids = torch.tensor([[0, 0, 0, 1, 1]]) - output = model(input_ids, token_type_ids=segment_ids, output_hidden_states=True) + with torch.no_grad(): + output = model(input_ids, token_type_ids=segment_ids, output_hidden_states=True) hidden_size = 1152