Merge pull request #2300 from GiulioZizzo/hf_model_wrapper_update

Huggingface model wrapper update

Author: beat-buesser

art/estimators/classification/hugging_face.py

@@ -156,6 +156,12 @@ def _make_model_wrapper(self, model: "torch.nn.Module") -> "torch.nn.Module":
 
         input_shape = self._input_shape
         input_for_hook = torch.rand(input_shape)
+        # self.device may not match the device the raw model was passed into ART.
+        # Check if the model is on cuda, if so set the hook input accordingly
+        if next(model.parameters()).is_cuda:
+            cuda_idx = torch.cuda.current_device()
+            input_for_hook = input_for_hook.to(torch.device(f"cuda:{cuda_idx}"))
+
         input_for_hook = torch.unsqueeze(input_for_hook, dim=0)
 
         if self.processor is not None: