No saving of checkpoints during last prediction phase of training after all epochs are saved

[paul-reiners]

We're running into a small problem. After the 1000 epochs of training are done, one final part is run, where predictions are made.

It looks like this:

0mo_SynthSeg_generated_0015 (3, 182, 218, 182)
debug: mirroring True mirror_axes (0, 1, 2)
step_size: 0.5
do mirror: True
data shape: (2, 182, 218, 182)
patch size: [112 128  96]
steps (x, y, and z): [[0, 35, 70], [0, 45, 90], [0, 43, 86]]
number of tiles: 27
computing Gaussian
prediction done

No checkpoints are saved during this section. So, if our SLURM script times out, nnU-Net starts over with this section when we start it up again (even with the -c option). Given a 24-hour SLURM time-out, It doesn't finish this section (for any of the folds).

Can we enable the saving of checkpoints of nnU-Net during this section?