fix max_epochs setup in basic example (Lightning-AI#1105)

tested only for the CPU version
tullie · Apr 3, 2020 · d0f6f71 · d0f6f71
1 parent 81b24cc
commit d0f6f71
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/pl_examples/basic_examples/cpu_template.py b/pl_examples/basic_examples/cpu_template.py
@@ -28,7 +28,7 @@ def main(hparams):
     # ------------------------
     # 2 INIT TRAINER
     # ------------------------
-    trainer = pl.Trainer()
+    trainer = pl.Trainer(max_epochs=hparams.epochs)
 
     # ------------------------
     # 3 START TRAINING

diff --git a/pl_examples/basic_examples/gpu_template.py b/pl_examples/basic_examples/gpu_template.py
@@ -29,6 +29,7 @@ def main(hparams):
     # 2 INIT TRAINER
     # ------------------------
     trainer = pl.Trainer(
+        max_epochs=hparams.epochs,
         gpus=hparams.gpus,
         distributed_backend=hparams.distributed_backend,
         use_amp=hparams.use_16bit