docstr: Test _set_module_training_mode
docstr_tok: ["Test","_set_module_training_mode"]
code: def test_set_module_training_mode(self) -> None:
      """
      Test _set_module_training_mode
      """
      module = nn.Linear(1, 1)
      loss_fn = nn.CrossEntropyLoss()
  
      tracked_modules: Dict[str, torch.nn.Module] = {
          "module": module,
          "loss_fn": loss_fn,
      }
  
      # set module training mode to False
      prior_module_train_states = _set_module_training_mode(
          tracked_modules, False
      )
  
      self.assertFalse(module.training)
      self.assertFalse(loss_fn.training)
  
      self.assertTrue(prior_module_train_states["module"])
      self.assertTrue(prior_module_train_states["loss_fn"])
  
      # set back to True
      prior_module_train_states = _set_module_training_mode(
          tracked_modules, True
      )
  
      self.assertTrue(module.training)
      self.assertTrue(loss_fn.training)
  
      self.assertFalse(prior_module_train_states["module"])
      self.assertFalse(prior_module_train_states["loss_fn"])
  
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 3 for calling the LoopUtilsTest obj's underlying member method code with example usage: obj.test_set_module_training_mode() without return types