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

2022-10-11 00:33:24 +05:30
parent d739a707d9
commit c6a928cadb
1 changed files with 2 additions and 1 deletions
--- a/tests/models/imagegpt/test_modeling_imagegpt.py
+++ b/tests/models/imagegpt/test_modeling_imagegpt.py
@@ -538,7 +538,8 @@ class ImageGPTModelIntegrationTest(unittest.TestCase):
        inputs = feature_extractor(images=image, return_tensors="pt").to(torch_device)

        # forward pass
-        outputs = model(**inputs)
+        with torch.no_grad():
+            outputs = model(**inputs)

        # verify the logits
        expected_shape = torch.Size((1, 1024, 512))