Update

1 files changed, 3 insertions, 1 deletions

diff --git a/train_lora.py b/train_lora.py
index 59beb09..f74a438 100644
--- a/train_lora.py
+++ b/train_lora.py
@@ -643,7 +643,9 @@ def main():
     num_train_epochs = args.num_train_epochs
     sample_frequency = args.sample_frequency
     if num_train_epochs is None:
-        num_train_epochs = math.ceil(args.num_train_steps / len(datamodule.train_dataset))
+        num_train_epochs = math.ceil(
+            args.num_train_steps / len(datamodule.train_dataset)
+        ) * args.gradient_accumulation_steps
         sample_frequency = math.ceil(num_train_epochs * (sample_frequency / args.num_train_steps))
 
     optimizer = create_optimizer(