docstr: Test _maybe_set_distributed_sampler_epoch util function
docstr_tok: ["Test","_maybe_set_distributed_sampler_epoch","util","function"]
code: def _test_maybe_set_distributed_sampler_epoch() -> None:
      """
      Test _maybe_set_distributed_sampler_epoch util function
      """
      dist.init_process_group("gloo")
  
      random_dataset = generate_random_dataset(10, 3)
      dummy_dataloader_with_distributed_sampler = DataLoader(
          random_dataset, sampler=DistributedSampler(random_dataset)
      )
  
      _maybe_set_distributed_sampler_epoch(
          dummy_dataloader_with_distributed_sampler, 20
      )
  
      sampler = cast(
          DistributedSampler[object],
          dummy_dataloader_with_distributed_sampler.sampler,
      )
      assert sampler.epoch == 20
  
imports: import unittest
  from typing import cast, Dict, Iterator
  import torch
  from torch import distributed, nn
  from torch.distributed import launcher
  from torch.utils.data import DataLoader
  from torch.utils.data.distributed import DistributedSampler
  from torchtnt.framework._loop_utils import _is_done, _is_epoch_done, _maybe_set_distributed_sampler_epoch, _reset_module_training_mode, _set_module_training_mode, _step_requires_iterator
  from torchtnt.framework._test_utils import generate_random_dataset
  from torchtnt.framework.state import State
  from torchtnt.utils.progress import Progress
  from torchtnt.utils.test_utils import get_pet_launch_config
label_desc: Use image node_id 2 for calling the LoopUtilsTest obj's underlying member method code with example usage: obj._test_maybe_set_distributed_sampler_epoch() without return types