Having validation sample plots indexed by epochs

[Sukh-P]

Just wanted to discuss whether it would be useful to try and get the slider/index for validation sample plots to which epoch it related to? (if this is possible)

e.g. we currently have:

Where the slider index is based on how many times the wandb.log() has been called (I think this is default behaviour) but it's a bit fiddly then to align which samples have been created which which epoch version of the model

[AUdaltsova]

I've looked into this briefly, it seems like wandb have added the tooling last year to make this possible in the UI (see wandb docs here). I couldn't figure out how to get it to work though, it kept complaining that my plots don't have epoch keys. I wonder if there is some different way to log the images to get them to have the epoch keys, although the docs suggest that run.log should be enough, which I think is equivalent to self.logger.experiment.log, which is what we use currently...