Fix CI: test_inference_for_pretraining in ViTMAEModelTest (#16591)
Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
This commit is contained in:
@@ -561,7 +561,7 @@ class ViTMAEModelIntegrationTest(unittest.TestCase):
|
|||||||
|
|
||||||
# forward pass
|
# forward pass
|
||||||
with torch.no_grad():
|
with torch.no_grad():
|
||||||
outputs = model(**inputs, noise=torch.from_numpy(noise))
|
outputs = model(**inputs, noise=torch.from_numpy(noise).to(device=torch_device))
|
||||||
|
|
||||||
# verify the logits
|
# verify the logits
|
||||||
expected_shape = torch.Size((1, 196, 768))
|
expected_shape = torch.Size((1, 196, 768))
|
||||||
|
|||||||
Reference in New Issue
Block a user