From 6eaf8bea693975eb772ca9a3749e0fc4b7fa48e4 Mon Sep 17 00:00:00 2001
From: Claire Vlases <99672427+cvlases@users.noreply.github.com>
Date: Fri, 5 Jun 2026 13:42:26 -0400
Subject: [PATCH] transfer to device not cpu

Move mean and std tensors to the specified device and update images_tensor accordingly.
---
 saliency_nodes.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/saliency_nodes.py b/saliency_nodes.py
index 06e0154..b4d98e7 100644
--- a/saliency_nodes.py
+++ b/saliency_nodes.py
@@ -132,8 +132,9 @@ def _make_call_model_fn(model, class_idx):
                 images_np = np.array(images, dtype=np.float32)
                 images_tensor = torch.from_numpy(images_np).permute(0, 3, 1, 2)
 
-                mean = torch.tensor(IMAGENET_MEAN).view(1, 3, 1, 1)
-                std  = torch.tensor(IMAGENET_STD).view(1, 3, 1, 1)
+                mmean = torch.tensor(IMAGENET_MEAN).view(1, 3, 1, 1).to(DEVICE)
+                std  = torch.tensor(IMAGENET_STD).view(1, 3, 1, 1).to(DEVICE)
+                images_tensor = images_tensor.to(DEVICE)
                 images_tensor = (images_tensor - mean) / std
 
                 # .clone() is required when escaping inference_mode