Lora fix: Save config JSON, too

1 files changed, 3 insertions, 0 deletions

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_