code: def wrap_model(
      model: nn.Module, grad_sample_mode: str, *args, **kwargs
  ):
      cls = get_gsm_class(grad_sample_mode)
      if grad_sample_mode == "functorch":
          kwargs["force_functorch"] = True
      return cls(model, *args, **kwargs)
  
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 2 for calling a global function with example usage: wrap_model(model, grad_sample_mode) and returns: cls