From 765bafb8e4ed11073f196a3372d7990428e270b0 Mon Sep 17 00:00:00 2001 From: Yih-Dar <2521628+ydshieh@users.noreply.github.com> Date: Tue, 5 Apr 2022 10:00:03 +0200 Subject: [PATCH] Fix CI: test_inference_for_pretraining in ViTMAEModelTest (#16591) Co-authored-by: ydshieh --- tests/vit_mae/test_modeling_vit_mae.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/vit_mae/test_modeling_vit_mae.py b/tests/vit_mae/test_modeling_vit_mae.py index 6ae62cb1c2..8cbde5b2ce 100644 --- a/tests/vit_mae/test_modeling_vit_mae.py +++ b/tests/vit_mae/test_modeling_vit_mae.py @@ -561,7 +561,7 @@ class ViTMAEModelIntegrationTest(unittest.TestCase): # forward pass 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 expected_shape = torch.Size((1, 196, 768))