Set torch.utils.data.DataLoader.batch_sampler epoch if defined

[Ghelfi]

🚀 Feature Request

Current implementation of composer trainer call the DistributedSampler.set_epoch method only on the Dataloader.sampler attribute but not on the Dataloader.batch_sampler, even if it is defined.
One example here.

Motivation

When doing distributed learning based on batch sampler one might want the epoch to be properly set on the batch_sampler since it is usually used to seed new seed over time.
This could be useful for Metric learning where batch_sampler can be a worth feature.
For now composer trainer only handle regular sampler.

Implementation

I'll propose a PR with a technical implementation. The idea is to check if batch_sampler is defined, in which case we can set the epoch properly, else we do it on the regular sampler (always defined in torch.utils.data.DalaLoader).

[Ghelfi]

The proposed PR: #3124

[Ghelfi]

Merged by @mvpatel2000. Closing the issue.