set_distributed_sampler_epoch for multidataset samplers (#1040)

Summary:
Pull Request resolved: #1040

- TorchTNT automatically sets the `sampler` epoch to ensure sampling is not the same across batches.
- However, the epoch was *not* being set for `batch_sampler`s, which are used in the `MultiDataset` case. See [a related (since resolved) issue for Pytorch Lightning](Lightning-AI/pytorch-lightning#13316).
- Here I'm running `set_epoch` for the `batch_sampler` when applicable.

Reviewed By: galrotem

Differential Revision: D87585393

fbshipit-source-id: e629253b45418710cf8c0bba118af24e37e2633e

Author: Rick Warren

torchtnt/framework/_loop_utils.py

@@ -77,11 +77,11 @@ def _maybe_set_distributed_sampler_epoch(
     See: https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler
     """
     # Set current training epoch for any DistributedSampler in dataloader
-    if isinstance(dataloader, torch.utils.data.DataLoader) and isinstance(
-        dataloader.sampler,
-        _DistributedSampler,
-    ):
-        dataloader.sampler.set_epoch(current_epoch)
+    if isinstance(dataloader, torch.utils.data.DataLoader):
+        if isinstance(dataloader.sampler, _DistributedSampler):
+            dataloader.sampler.set_epoch(current_epoch)
+        elif isinstance(dataloader.batch_sampler, _DistributedSampler):
+            dataloader.batch_sampler.set_epoch(current_epoch)
 
 
 def _set_module_training_mode(