Added max_steps parameter to set the maximum number of steps to train

jshuadvd · Jul 9, 2024 · 6b32640 · 6b32640
1 parent f2715ae
commit 6b32640
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/train.py b/train.py
@@ -140,6 +140,7 @@ def train(
     epochs=10,
     gradient_accumulation_steps=4,
     resume_from_checkpoint=None,
+    max_steps=None,
 ):
     """
     Train the LongRoPE model.
@@ -155,6 +156,7 @@ def train(
         epochs (int): Number of training epochs.
         gradient_accumulation_steps (int): Number of steps to accumulate gradients.
         resume_from_checkpoint (str): Path to a checkpoint to resume training from.
+        max_steps (int): Maximum number of steps to train. If None, train for full epochs.
 
     Returns:
         None