[Tests] Fix DiT test (#16218)

* Fix device

* Clean up

Co-authored-by: Niels Rogge <nielsrogge@Nielss-MacBook-Pro.local>

tests/dit/test_modeling_dit.py

@@ -43,7 +43,7 @@ class DiTIntegrationTest(unittest.TestCase):
 
         image = dataset["train"][0]["image"].convert("RGB")
 
-        inputs = feature_extractor(image, return_tensors="pt")
+        inputs = feature_extractor(image, return_tensors="pt").to(torch_device)
 
         # forward pass
         with torch.no_grad():