From a1946f09013ab95e35fb47cfec17d820bf5ec6b2 Mon Sep 17 00:00:00 2001
From: Claire Vlases <99672427+cvlases@users.noreply.github.com>
Date: Fri, 5 Jun 2026 13:31:33 -0400
Subject: [PATCH] Update model evaluation to use gpu for on the cloud

Change model evaluation to use available device (CPU or CUDA).
---
 saliency_nodes.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/saliency_nodes.py b/saliency_nodes.py
index ca48c86..6f87916 100644
--- a/saliency_nodes.py
+++ b/saliency_nodes.py
@@ -79,7 +79,8 @@ def _load_model(model_name: str):
                 except AttributeError:
                     model = constructor(pretrained=True)
 
-                model.eval().cpu()
+                device = "cuda" if torch.cuda.is_available() else "cpu"
+                model.eval().to(device)
                 # Deep clone all params and buffers to ensure no inference tensors remain
                 for param in model.parameters():
                     param.data = param.data.clone()
@@ -490,4 +491,4 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "SaliencyMap":        "Saliency Map (PAIR)",
     "SaliencyComparison": "Saliency Comparison Grid (PAIR)",
     "SaliencyArt":        "Explainable AI",
-}
\ No newline at end of file
+}