1 files changed, 11 insertions, 2 deletions

diff --git a/train_lora.py b/train_lora.py
index c197206..9cf17c7 100644
--- a/train_lora.py
+++ b/train_lora.py
@@ -23,7 +23,7 @@ from data.csv import VlpnDataModule, keyword_filter
 from training.functional import train, add_placeholder_tokens, get_models
 from training.strategy.lora import lora_strategy
 from training.optimization import get_scheduler
-from training.util import save_args
+from training.util import AverageMeter, save_args
 
 # https://github.com/huggingface/peft/blob/main/examples/lora_dreambooth/train_dreambooth.py
 UNET_TARGET_MODULES = ["to_q", "to_v", "query", "value"]
@@ -1035,6 +1035,11 @@ def main():
     lr_warmup_epochs = args.lr_warmup_epochs
     lr_cycles = args.lr_cycles
 
+    avg_loss = AverageMeter()
+    avg_acc = AverageMeter()
+    avg_loss_val = AverageMeter()
+    avg_acc_val = AverageMeter()
+
     while True:
         if len(auto_cycles) != 0:
             response = auto_cycles.pop(0)
@@ -1122,7 +1127,7 @@ def main():
             warmup_epochs=lr_warmup_epochs,
         )
 
-        lora_checkpoint_output_dir = output_dir / lora_project / f"model_{training_iter + 1}"
+        lora_checkpoint_output_dir = output_dir / lora_project / f"model_{training_iter}"
 
         trainer(
             strategy=lora_strategy,
@@ -1142,6 +1147,10 @@ def main():
             sample_frequency=lora_sample_frequency,
             offset_noise_strength=args.offset_noise_strength,
             no_val=args.valid_set_size == 0,
+            avg_loss=avg_loss,
+            avg_acc=avg_acc,
+            avg_loss_val=avg_loss_val,
+            avg_acc_val=avg_acc_val,
         )
 
         training_iter += 1