From e5972967a860d551a03ae598990b2252f0387c56 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Thu, 23 May 2024 15:36:21 +0200 Subject: [PATCH] fix on_train_end not being executed in OOM cases --- .../variants/benchmarking/nnUNetTrainerBenchmark_5epochs.py | 1 + 1 file changed, 1 insertion(+) diff --git a/nnunetv2/training/nnUNetTrainer/variants/benchmarking/nnUNetTrainerBenchmark_5epochs.py b/nnunetv2/training/nnUNetTrainer/variants/benchmarking/nnUNetTrainerBenchmark_5epochs.py index a76bb634c..f11e44e01 100644 --- a/nnunetv2/training/nnUNetTrainer/variants/benchmarking/nnUNetTrainerBenchmark_5epochs.py +++ b/nnunetv2/training/nnUNetTrainer/variants/benchmarking/nnUNetTrainerBenchmark_5epochs.py @@ -29,6 +29,7 @@ def run_training(self): super().run_training() except RuntimeError: self.crashed_with_runtime_error = True + self.on_train_end() def on_train_end(self): super().on_train_end()