From aba41e4d8b395d65e1b927367c19f02f19d80240 Mon Sep 17 00:00:00 2001
From: Bruno Fargnoli <visualbruno@live.fr>
Date: Sat, 20 Dec 2025 11:05:56 +0100
Subject: [PATCH] Updated "Load Model" node

---
 nodes.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/nodes.py b/nodes.py
index 54bb3ab..affc416 100644
--- a/nodes.py
+++ b/nodes.py
@@ -114,7 +114,14 @@ class Trellis2LoadModel:
         
         pipeline = Trellis2ImageTo3DPipeline.from_pretrained(model_path)
         pipeline.low_vram = low_vram
-        pipeline.to(device)
+        
+        if device=="cuda":
+            if low_vram:
+                pipeline.cuda()
+            else:
+                pipeline.to(device)
+        else:
+            pipeline.to(device)
         
         return (pipeline,)