Inspect the model's media when the model improved

[matekrk]

I am not sure if it is a question or request for new feature.

I use Trainer framework with ModelCheckpoint (let's say k=2). During the training, my models multiple times get better, however I would like to log (eg. using wandb) the images and visually see the progress of the image quality with each improvement without logging in every epoch or without saving all intermediate models.

Motivation: during the training (especially when using the warm restarts), multiple times we may encounter the improvement on the (validation) metrics. However, when working with images, sometimes it is beneficial to not only rely on the single number such as MSE, but also see the visual improvement.

I tried so far with overwriting save_checkpoint function, but it did not work as desired.