Custom batch selection for logging

[bhosalems]

Description & Motivation

Need to be able to select the same batch in every logging cycle. For generation pipelines similar to stable diffusion it is very hard to gauge the performance over training if we continue to choose random batches.

Pitch

User should have selective ability to choose the batch to log which will be constant for all the logging cycles.

Alternatives

Its possible to load the data again in train_btach_end() or validation_batch_end(), and call logging.

Additional context

No response

cc @borda