fix : save lora 16 precision (#97)

cloneofsimo · Dec 29, 2022 · 35653d2 · 35653d2
1 parent 5240d32
commit 35653d2
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/lora_diffusion/lora.py b/lora_diffusion/lora.py
@@ -185,8 +185,8 @@ def save_lora_weight(
     for _up, _down in extract_lora_ups_down(
         model, target_replace_module=target_replace_module
     ):
-        weights.append(_up.weight)
-        weights.append(_down.weight)
+        weights.append(_up.weight.to("cpu").to(torch.float16))
+        weights.append(_down.weight.to("cpu").to(torch.float16))
 
     torch.save(weights, path)