From 28007bd0764daf91208452420caabbe9c7055905 Mon Sep 17 00:00:00 2001 From: jere357 Date: Mon, 3 Jul 2023 14:26:50 +0200 Subject: [PATCH] Consistent early exit logic for validation loop in Fabric Trainer example #17962 --- examples/fabric/build_your_own_trainer/trainer.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/examples/fabric/build_your_own_trainer/trainer.py b/examples/fabric/build_your_own_trainer/trainer.py index 492245e699..08692377d8 100644 --- a/examples/fabric/build_your_own_trainer/trainer.py +++ b/examples/fabric/build_your_own_trainer/trainer.py @@ -216,8 +216,7 @@ class MyCustomTrainer: for batch_idx, batch in enumerate(iterable): # end epoch if stopping training completely or max batches for this epoch reached if self.should_stop or batch_idx >= limit_batches: - self.fabric.call("on_train_epoch_end") - return + break self.fabric.call("on_train_batch_start", batch, batch_idx)