1 files changed, 5 insertions, 18 deletions

diff --git a/training/strategy/lora.py b/training/strategy/lora.py
index 92abaa6..bc10e58 100644
--- a/training/strategy/lora.py
+++ b/training/strategy/lora.py
@@ -89,20 +89,14 @@ def lora_strategy_callbacks(
     @torch.no_grad()
     def on_checkpoint(step, postfix):
         print(f"Saving checkpoint for step {step}...")
-        orig_unet_dtype = unet.dtype
-        unet.to(dtype=torch.float32)
-        unet.save_attn_procs(checkpoint_output_dir.joinpath(f"{step}_{postfix}"))
-        unet.to(dtype=orig_unet_dtype)
+
+        unet_ = accelerator.unwrap_model(unet)
+        unet_.save_attn_procs(checkpoint_output_dir.joinpath(f"{step}_{postfix}"))
+        del unet_
 
     @torch.no_grad()
     def on_sample(step):
-        orig_unet_dtype = unet.dtype
-        unet.to(dtype=weight_dtype)
         save_samples_(step=step)
-        unet.to(dtype=orig_unet_dtype)
-
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
 
     return TrainingCallbacks(
         on_prepare=on_prepare,
@@ -126,16 +120,9 @@ def lora_prepare(
     lora_layers: AttnProcsLayers,
     **kwargs
 ):
-    weight_dtype = torch.float32
-    if accelerator.state.mixed_precision == "fp16":
-        weight_dtype = torch.float16
-    elif accelerator.state.mixed_precision == "bf16":
-        weight_dtype = torch.bfloat16
-
     lora_layers, optimizer, train_dataloader, val_dataloader, lr_scheduler = accelerator.prepare(
         lora_layers, optimizer, train_dataloader, val_dataloader, lr_scheduler)
-    unet.to(accelerator.device, dtype=weight_dtype)
-    text_encoder.to(accelerator.device, dtype=weight_dtype)
+
     return text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler, {"lora_layers": lora_layers}