Merge branch 'dev_1.17.0' into composite-adversarial-attack

Trusted-AI · Dec 16, 2023 · d8f2a10 · d8f2a10
2 parents c8f5816 + dec5ddf
commit d8f2a10
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 17 deletions.
diff --git a/art/estimators/classification/hugging_face.py b/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:

diff --git a/notebooks/huggingface_notebook.ipynb b/notebooks/huggingface_notebook.ipynb