From d97ba098855783fca9124b836c96d8f0d9ba8e03 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sat, 10 Feb 2024 13:05:40 +0200
Subject: [PATCH] Update nodes.py

---
 nodes.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/nodes.py b/nodes.py
index 7c46542..7835ca0 100644
--- a/nodes.py
+++ b/nodes.py
@@ -770,8 +770,9 @@ class TrainMotionDirectorLora:
 
             use_offset_noise = False
 
-            device = unet.device
-
+            device = comfy.model_management.get_torch_device()
+            unet.to(device)
+            
             target_spatial_modules = ["Transformer3DModel"]
             target_temporal_modules = ["TemporalTransformerBlock"]
 
@@ -784,7 +785,7 @@ class TrainMotionDirectorLora:
 
             # Support mixed-precision training
             scaler = torch.cuda.amp.GradScaler()
-
+            
             num_update_steps_per_epoch = math.ceil(batch_size) / gradient_accumulation_steps
             num_train_epochs = math.ceil(max_train_steps / num_update_steps_per_epoch)