Call `set_epoch` on `Dataloader.batch_sampler` if defined

[Ghelfi]

What does this PR do?

Proposes a technical implementation to set_epoch on either DataLoader.sampler or DataLoader.batch_sampler if it is defined.

What issue(s) does this change relate to?

• Related to #3123

[mvpatel2000]

Thanks for the PR!

Would you mind adding a unit test too please? You can extend test_checkpoint.py with class TestCheckpointResumption:. I would add an arg to get_trainer to add a batch dist sampler and add a new fn that tests that path instead.

If you need help, I'm happy to provide assistance for this!

When you are ready, please request review from mvpatel2000!

[Ghelfi]

@mvpatel2000 cannot add you as reviewer to ask for a review. I added a test and ready for a review. I put the test where you suggested. Will it be better to put it in test_trainer.py?

[mvpatel2000]

@Ghelfi do you mind if I push directly to this branch? there's a few lint things and refactor in test I would prefer. I'm happy to leave as comments if you'd like, but I figure it's faster for me to directly apply. I don't want to delay you further :)

[Ghelfi]

@mvpatel2000 fine by me. Thanks for handling it.