1 files changed, 2 insertions, 3 deletions

diff --git a/training/strategy/ti.py b/training/strategy/ti.py
index 6a76f98..14bdafd 100644
--- a/training/strategy/ti.py
+++ b/training/strategy/ti.py
@@ -176,10 +176,9 @@ def textual_inversion_prepare(
     elif accelerator.state.mixed_precision == "bf16":
         weight_dtype = torch.bfloat16
 
-    prep = [text_encoder] + list(args)
-    text_encoder, optimizer, train_dataloader, val_dataloader, lr_scheduler = accelerator.prepare(*prep)
+    prepped = accelerator.prepare(text_encoder, *args)
     unet.to(accelerator.device, dtype=weight_dtype)
-    return text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler
+    return (prepped[0], unet) + prepped[1:]
 
 
 textual_inversion_strategy = TrainingStrategy(