class_docstr: Extends nn.Module so that its parameter tensors have an extra field called .grad_sample.
class_docstr_tok: ["Extends","nn.Module","so","that","its","parameter","tensors","have","an","extra","field","called",".grad_sample","."]
docstr: Args:
      m: nn.Module to be wrapped
      batch_first: Flag to indicate if the input tensor to the corresponding module
          has the first dimension representing the batch. If set to True, dimensions on
          input tensor are expected be ``[batch_size, ...]``, otherwise
          ``[K, batch_size, ...]``
      loss_reduction: Indicates if the loss reduction (for aggregating the gradients)
          is a sum or a mean operation. Can take values "sum" or "mean"
  
  Raises:
      NotImplementedError
          If ``strict`` is set to ``True`` and module ``m`` (or any of its
          submodules) doesn't have a registered grad sampler function.
docstr_tok: ["Args",":","m",":","nn.Module","to","be","wrapped","batch_first",":","Flag","to","indicate","if","the","input","tensor","to","the","corresponding","module","has","the","first","dimension","representing","the","batch",".","If","set","to","True",",","dimensions","on","input","tensor","are","expected","be","``","[","batch_size",",","...","]","``",",","otherwise","``","[","K",",","batch_size",",","...","]","``","loss_reduction",":","Indicates","if","the","loss","reduction","(","for","aggregating","the","gradients",")","is","a","sum","or","a","mean","operation",".","Can","take","values","``","sum","''","or","``","mean","''","Raises",":","NotImplementedError","If","``","strict","``","is","set","to","``","True","``","and","module","``","m","``","(","or","any","of","its","submodules",")","does","n't","have","a","registered","grad","sampler","function","."]
code: def __init__(
      self,
      m: nn.Module,
      *,
      batch_first=True,
      loss_reduction="mean",
  ):
      """
  
      Args:
          m: nn.Module to be wrapped
          batch_first: Flag to indicate if the input tensor to the corresponding module
              has the first dimension representing the batch. If set to True, dimensions on
              input tensor are expected be ``[batch_size, ...]``, otherwise
              ``[K, batch_size, ...]``
          loss_reduction: Indicates if the loss reduction (for aggregating the gradients)
              is a sum or a mean operation. Can take values "sum" or "mean"
  
      Raises:
          NotImplementedError
              If ``strict`` is set to ``True`` and module ``m`` (or any of its
              submodules) doesn't have a registered grad sampler function.
      """
      super().__init__()
  
      self._module = m
      self.batch_first = batch_first
      self.loss_reduction = loss_reduction
      self.grad_accumulation_hook: Optional[RemovableHandle] = None
  
      for _, p in trainable_parameters(self):
          p.grad_sample = None
          p._forward_counter = 0
  
imports: import logging
  from abc import ABC, abstractmethod
  from typing import Optional
  import torch.nn
  from opacus.utils.module_utils import trainable_parameters
  from torch.utils.hooks import RemovableHandle
label_desc: Use image node_id 1 to create a new AbstractGradSampleModule object from inherited base classes: nn, ABC with example: obj = AbstractGradSampleModule(m)