code: def get_gsm_class(
      grad_sample_mode: str,
  ) -> Type[AbstractGradSampleModule]:
      """
      Returns AbstractGradSampleModule subclass correspinding to the input mode.
      See README for detailed comparison between grad sample modes.
  
      :param grad_sample_mode:
      :return:
      """
      if grad_sample_mode in ["hooks", "functorch"]:
          return GradSampleModule
      elif grad_sample_mode == "ew":
          return GradSampleModuleExpandedWeights
      elif grad_sample_mode == "no_op":
          return GradSampleModuleNoOp
      else:
          raise ValueError(
              f"Unexpected grad_sample_mode: {grad_sample_mode}. "
              f"Allowed values: hooks, ew"
          )
  
file_path: opacus/opacus/grad_sample/utils.py
imports: from typing import Sequence, Type, Union
  import torch.nn
  from .grad_sample_module import GradSampleModule
  from .gsm_base import AbstractGradSampleModule
  from .gsm_exp_weights import GradSampleModuleExpandedWeights
  from .gsm_no_op import GradSampleModuleNoOp
label_desc: Use image node_id 3 for calling a global function with example usage: get_gsm_class(grad_sample_mode) and returns: GradSampleModule, GradSampleModuleExpandedWeights, GradSampleModuleNoOp