Set dtype and device for pixel_values in forward

huggingface · Nov 4, 2024 · 9767048 · 9767048
1 parent ec9eade
commit 9767048
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/transformers/models/timm_wrapper/modeling_timm_wrapper.py b/src/transformers/models/timm_wrapper/modeling_timm_wrapper.py
@@ -112,6 +112,7 @@ def forward(
         if output_hidden_states is not None or output_attentions is not None:
             raise ValueError("Cannot set output_attentions or output_hidden_states for timm models")
 
+        pixel_values = pixel_values.to(self.device, self.dtype)
         features = self.timm_model(pixel_values, **kwargs)
 
         if not return_dict:
@@ -152,6 +153,7 @@ def forward(
         if output_hidden_states is not None or output_attentions is not None:
             raise ValueError("Cannot set `output_attentions` or `output_hidden_states` for timm models")
 
+        pixel_values = pixel_values.to(self.device, self.dtype)
         logits = self.timm_model(pixel_values, **kwargs)
 
         loss = None