add maxiters to trainer

gaspreth · Jul 1, 2022 · d9ea878 · d9ea878
1 parent 1f33571
commit d9ea878
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/mingpt/trainer.py b/mingpt/trainer.py
@@ -20,6 +20,7 @@ def get_default_config(cls):
         # dataloder parameters
         C.num_workers = 4
         # optimizer parameters
+        C.max_iters = None
         C.batch_size = 64
         C.learning_rate = 3e-4
         C.betas = (0.9, 0.95)
@@ -99,3 +100,6 @@ def run(self):
             self.iter_dt = tnow - self.iter_time
             self.iter_time = tnow
 
+            # termination conditions
+            if config.max_iters is not None and self.iter_num >= config.max_iters:
+                break