Update nodes.py

2024-02-10 14:57:04 +02:00
parent f070cc0e97
commit 2911a53a9a
1 changed files with 2 additions and 0 deletions
@@ -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"]