2 files changed, 3 insertions, 6 deletions

diff --git a/training/strategy/dreambooth.py b/training/strategy/dreambooth.py
index 1277939..e88bf90 100644
--- a/training/strategy/dreambooth.py
+++ b/training/strategy/dreambooth.py
@@ -193,9 +193,7 @@ def dreambooth_prepare(
     unet: UNet2DConditionModel,
     *args
 ):
-    prep = [text_encoder, unet] + list(args)
-    text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler = accelerator.prepare(*prep)
-    return text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler
+    return accelerator.prepare(text_encoder, unet, *args)
 
 
 dreambooth_strategy = TrainingStrategy(
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(