code: def register_grad_sampler(
      target_class_or_classes: Union[
          Type[nn.Module], Sequence[Type[nn.Module]]
      ]
  ):
      """
      Registers the decorated function as the ``grad_sampler`` of ``target_class_or_classes``, which is
      the function that will be invoked every time you want to compute a per-sample gradient
      of ``target_class_or_classes``. The signature of every grad_sampler is always the same:
  
      >>> @register_grad_sampler(MyCustomModel)
      ... def compute_grad_sample(module, activations, backprops):
      ...    pass
  
      It may help you to take a look at the existing grad_samplers inside Opacus, under ``opacus.grad_sample.``
      """
  
      def decorator(f):
          target_classes = (
              target_class_or_classes
              if isinstance(target_class_or_classes, Sequence)
              else [target_class_or_classes]
          )
          for target_class in target_classes:
              GradSampleModule.GRAD_SAMPLERS[target_class] = f
          return f
  
      return decorator
  
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 1 for calling a global function with example usage: register_grad_sampler(target_class_or_classes) and returns: decorator, f