Unverified Commit 765bafb8 authored by Yih-Dar's avatar Yih-Dar Committed by GitHub
Browse files

Fix CI: test_inference_for_pretraining in ViTMAEModelTest (#16591)


Co-authored-by: default avatarydshieh <ydshieh@users.noreply.github.com>
parent 104c0652
...@@ -561,7 +561,7 @@ class ViTMAEModelIntegrationTest(unittest.TestCase): ...@@ -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))
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment