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
+}