From 1383f64a5f1f55c39c2df47c6fd778d1870afee5 Mon Sep 17 00:00:00 2001 From: Ondrej Platek Date: Thu, 12 Mar 2020 17:42:52 +0100 Subject: [PATCH] fix max_epochs setup in basic example (#1105) tested only for the CPU version --- pl_examples/basic_examples/cpu_template.py | 2 +- pl_examples/basic_examples/gpu_template.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/pl_examples/basic_examples/cpu_template.py b/pl_examples/basic_examples/cpu_template.py index 0714b1aac0..7781ba3271 100644 --- 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 index c661eef65f..090b2adcba 100644 --- 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