fix: lora config saving with multiple gpus

jolibrain · May 18, 2024 · c98617d · c98617d
1 parent f7f26a0
commit c98617d
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/models/base_model.py b/models/base_model.py
@@ -801,7 +801,17 @@ def save_networks(self, epoch):
                 net = getattr(self, "net" + name)
 
                 if len(self.gpu_ids) > 1 and self.use_cuda:
-                    torch.save(net.module.state_dict(), save_path)
+                    if (
+                        name == "G_A"
+                        and hasattr(net.module, "unet")
+                        and hasattr(net.module, "vae")
+                        and any(
+                            "lora" in n for n, _ in net.module.unet.named_parameters()
+                        )
+                    ):
+                        net.module.save_lora_config(save_path)
+                    else:
+                        torch.save(net.module.state_dict(), save_path)
                 else:
                     if (
                         name == "G_A"