fix on_train_end not being executed in OOM cases

MIC-DKFZ · May 23, 2024 · e597296 · e597296
1 parent e0b33a6
commit e597296
Showing 1 changed file with 1 addition and 0 deletions.
diff --git 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()