From be3e05e47cded8487aaa787c54aa74770f9dcac8 Mon Sep 17 00:00:00 2001
From: Volpeon <git@volpeon.ink>
Date: Fri, 31 Mar 2023 09:34:55 +0200
Subject: Fix

---
 train_lora.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

(limited to 'train_lora.py')

diff --git a/train_lora.py b/train_lora.py
index 6f8644b..9975462 100644
--- a/train_lora.py
+++ b/train_lora.py
@@ -630,7 +630,8 @@ def main():
     num_train_epochs = args.num_train_epochs
 
     if num_train_epochs is None:
-        num_train_epochs = math.ceil(args.num_train_steps / len(datamodule.train_dataset))
+        num_images = math.ceil(len(datamodule.train_dataset) / args.train_batch_size) * args.train_batch_size
+        num_train_epochs = math.ceil(args.num_train_steps / num_images)
 
     optimizer = create_optimizer(
         itertools.chain(
-- 
cgit v1.2.3-54-g00ecf