Merge remote-tracking branch 'origin/main' into antonsc/fastmri

InnerEye/ML/model_training.py

@@ -117,13 +117,10 @@ def create_lightning_trainer(container: LightningContainer,
  save_last=True)
 
  # Recovery checkpoints: {epoch} will turn into a string like "epoch=1"
- # Store 1 recovery checkpoint every recovery_checkpoint_save_interval epochs. Due to a bug in Lightning, this
- # will still write alternate files recovery.ckpt and recovery-v0.ckpt, which are cleaned up later in
- # cleanup_checkpoint_folder
- recovery_checkpoint_callback = ModelCheckpoint(dirpath=str(container.checkpoint_folder),
- filename=RECOVERY_CHECKPOINT_FILE_NAME,
- period=container.recovery_checkpoint_save_interval
- )
+ # Store 1 recovery checkpoint every recovery_checkpoint_save_interval epochs, keep the last
+ # recovery_checkpoints_save_last_k.
+ recovery_checkpoint_callback = InnerEyeRecoveryCheckpointCallback(container)
+
  available_gpus = torch.cuda.device_count()
  num_gpus = available_gpus if container.use_gpu else 0
  no_gpu_message = "" if container.use_gpu else ". Not using any GPU because the use_gpu flag is set to False."