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

amyeroberts · Oct 18, 2022 · d3c0436 · d3c0436
1 parent b6c12ab
commit d3c0436
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/models/imagegpt/test_modeling_imagegpt.py b/tests/models/imagegpt/test_modeling_imagegpt.py
@@ -538,7 +538,8 @@ def test_inference_causal_lm_head(self):
         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))