move hook input to original model device

Signed-off-by: GiulioZizzo <[email protected]>

Author: GiulioZizzo

art/estimators/classification/hugging_face.py

@@ -155,7 +155,7 @@ def _make_model_wrapper(self, model: "torch.nn.Module") -> "torch.nn.Module":
         import torch
 
         input_shape = self._input_shape
-        input_for_hook = torch.rand(input_shape)
+        input_for_hook = torch.rand(input_shape).to(self.device)
         input_for_hook = torch.unsqueeze(input_for_hook, dim=0)
 
         if self.processor is not None: