code: def train_mode(*nets: nn.Module):
      """
      Set network(s) to train mode and then return to original state at the end.
  
      Args:
          nets: Input network(s)
  
      Examples
  
      .. code-block:: python
  
          t=torch.rand(1,1,16,16)
          p=torch.nn.Conv2d(1,1,3)
          p.eval()
          print(p.training)  # False
          with train_mode(p):
              print(p.training)  # True
              print(p(t).sum().backward())  # No exception
      """
  
      # Get original state of network(s)
      # Check the training attribute in case it's TensorRT based models which don't have this attribute.
      eval_list = [
          n for n in nets if hasattr(n, "training") and (not n.training)
      ]
  
      try:
          # set to train mode
          with torch.set_grad_enabled(True):
              yield [
                  n.train() if hasattr(n, "train") else n for n in nets
              ]
      finally:
          # Return required networks to eval_list
          for n in eval_list:
              if hasattr(n, "eval"):
                  n.eval()
  
imports: from __future__ import annotations
  import io
  import re
  import warnings
  from collections import OrderedDict
  from collections.abc import Callable, Mapping, Sequence
  from contextlib import contextmanager
  from copy import deepcopy
  from typing import Any
  import numpy
  import torch
  import torch.nn
  from monai.apps.utils import get_logger
  from monai.config import PathLike
  from monai.utils.misc import ensure_tuple, save_obj, set_determinism
  from monai.utils.module import look_up_option, optional_import, pytorch_after
  from monai.utils.type_conversion import convert_to_dst_type, convert_to_tensor
label_desc: Use image node_id 12 for calling a global function with example usage: train_mode() without return types