Lightning: set limit_train_batches and limit_val_batches (horovod#3237)

Tell Lightning trainer that how many batches a single epoch needs. Signed-off-by: Chongxiao Cao <chongxiaoc@uber.com> Signed-off-by: weihanmines <weihan13@amd.com>
weihanmines · Dec 10, 2021 · 0e61f3c · 0e61f3c
1 parent 8c2a307
commit 0e61f3c
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/horovod/spark/lightning/remote.py b/horovod/spark/lightning/remote.py
@@ -193,6 +193,8 @@ def on_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -
                       'gpus': _num_gpus,
                       'callbacks': callbacks,
                       'max_epochs': epochs,
+                      'limit_train_batches': _train_steps_per_epoch,
+                      'limit_val_batches': _val_steps_per_epoch,
                       'logger': train_logger,
                       'log_every_n_steps': log_every_n_steps,
                       'num_sanity_val_steps': 0,