class_docstr: ## Zero3 Layer
  
  Each layer of the model (or a combination of a few consecutive layers) should be wrapped in
  this module.
class_docstr_tok: ["#","#","Zero3","Layer","Each","layer","of","the","model","(","or","a","combination","of","a","few","consecutive","layers",")","should","be","wrapped","in","this","module","."]
docstr: #### Add backward hooks to the parameters of the current layer.
docstr_tok: ["#","#","#","#","Add","backward","hooks","to","the","parameters","of","the","current","layer","."]
code: def _add_backward_hooks(self):
      """
      #### Add backward hooks to the parameters of the current layer.
      """
  
      # Number of backward hooks added
      self._backward_hook_handles = 0
  
      # Loop through trainable parameters of the current layer
      for p in self.param_refs[self.TRAINING_PARAMS_IDX]:
          # Make sure a hook hasn't already been added
          assert not hasattr(
              p, "_hook_handle"
          ), "Parameter has already been hooked"
          # Use `expand_as` to create an autograd step which we can intercept
          p_tmp = p.expand_as(p)
          # Get a handle to add the backward hook.
          # [This blog discusses about `grad_acc`](https://amsword.medium.com/understanding-pytorchs-autograd-with-grad-fn-and-next-functions-b2c4836daa00).
          grad_acc = p_tmp.grad_fn.next_functions[0][0]
          # Add the backward hook
          handle = grad_acc.register_hook(
              functools.partial(self._post_backward_hook, p)
          )
          # Keep a reference to the handle
          p._hook_handle = handle
          # Increment the number of hooks added
          self._backward_hook_handles += 1
  
imports: import functools
  from typing import List, Optional, Tuple
  import torch
  import torch.distributed
  from torch import nn
label_desc: Use image node_id 8 for calling the Zero3Layer obj's underlying member method code with example usage: obj._add_backward_hooks() without return types