From 2911a53a9a81cfe2ee5dbeb279c8953e478f3e4e Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sat, 10 Feb 2024 14:57:04 +0200
Subject: [PATCH] Update nodes.py

---
 nodes.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/nodes.py b/nodes.py
index b8abb9c..bed0c04 100644
--- a/nodes.py
+++ b/nodes.py
@@ -751,6 +751,8 @@ class TrainMotionDirectorLora:
             device = comfy.model_management.get_torch_device()
             unet.to(device)
             vae.to(device)
+            text_encoder.to(device)
+            tokenizer.to(device)
 
             target_spatial_modules = ["Transformer3DModel"]
             target_temporal_modules = ["TemporalTransformerBlock"]