Update tests/models/textnet/test_modeling_textnet.py

Co-authored-by: Pavel Iakubovskii <[email protected]>
huggingface · Dec 21, 2024 · 5f2b970 · 5f2b970
1 parent ed36fef
commit 5f2b970
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/models/textnet/test_modeling_textnet.py b/tests/models/textnet/test_modeling_textnet.py
@@ -324,7 +324,8 @@ def test_inference_textnet_image_classification(self):
         inputs = processor(images=image, return_tensors="pt").to(torch_device)
 
         # forward pass
-        output = model(**inputs)
+        with torch.no_grad():
+            output = model(**inputs)
 
         # verify logits
         self.assertEqual(output.logits.shape, torch.Size([1, 2]))