Added forward pass of test_inference_image_classification_head with torch.no_grad() (#14777)
This commit is contained in:
@@ -347,7 +347,8 @@ class ViTModelIntegrationTest(unittest.TestCase):
|
|||||||
inputs = feature_extractor(images=image, return_tensors="pt").to(torch_device)
|
inputs = feature_extractor(images=image, return_tensors="pt").to(torch_device)
|
||||||
|
|
||||||
# forward pass
|
# forward pass
|
||||||
outputs = model(**inputs)
|
with torch.no_grad():
|
||||||
|
outputs = model(**inputs)
|
||||||
|
|
||||||
# verify the logits
|
# verify the logits
|
||||||
expected_shape = torch.Size((1, 1000))
|
expected_shape = torch.Size((1, 1000))
|
||||||
|
|||||||
Reference in New Issue
Block a user