From f0471f21f419a34e3fd7b7b03a4292c139fda674 Mon Sep 17 00:00:00 2001
From: Volpeon <git@volpeon.ink>
Date: Fri, 24 Mar 2023 10:58:10 +0100
Subject: Lora fix: Save config JSON, too

---
 training/strategy/lora.py | 3 +++
 1 file changed, 3 insertions(+)

(limited to 'training/strategy')

diff --git a/training/strategy/lora.py b/training/strategy/lora.py
index 3971eae..1e32114 100644
--- a/training/strategy/lora.py
+++ b/training/strategy/lora.py
@@ -3,6 +3,7 @@ from functools import partial
 from contextlib import contextmanager
 from pathlib import Path
 import itertools
+import json
 
 import torch
 from torch.utils.data import DataLoader
@@ -95,6 +96,8 @@ def lora_strategy_callbacks(
 
         accelerator.print(state_dict)
         accelerator.save(state_dict, checkpoint_output_dir / f"{step}_{postfix}.pt")
+        with open(checkpoint_output_dir / "lora_config.json", "w") as f:
+            json.dump(lora_config, f)
 
         del unet_
         del text_encoder_
-- 
cgit v1.2.3-70-g09d2