code: def _set_module_training_mode(
      modules: Dict[str, nn.Module], mode: bool
  ) -> Dict[str, bool]:
      """Returns states to allow for a reset at the end of the loop."""
      prior_module_train_states = {}
      for name, module in modules.items():
          prior_module_train_states[name] = module.training
          module.train(mode)
      return prior_module_train_states
  
file_path: tnt/torchtnt/framework/_loop_utils.py
imports: import collections
  import inspect
  import logging
  from typing import Callable, Dict, Iterable, Optional, TypeVar
  import torch
  import torch.nn
  import typing_extensions
  from torchtnt.framework.state import State
  from torchtnt.utils.progress import Progress
label_desc: Use image node_id 4 for calling a global function with example usage: _set_module_training_mode(modules, mode) and returns: prior_module_train_states