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: ### Forward pass
docstr_tok: ["#","#","#","Forward","pass"]
code: def forward(self, *args, **kwargs):
      """
      ### Forward pass
      """
  
      # Fetch all the parameters of the current node.
      # This gets called by the previous layer so this call is just to make sure parameters are fetched.
      self.fetch_params()
  
      # Wait for parameter fetching to complete.
      torch.cuda.current_stream().wait_stream(self.fetch_stream)
  
      # Start fetching parameters of the proceeding layers, so that they will fetch them which the current layer
      # does its computations.
      for layer in self.next_layer:
          layer.fetch_params()
  
      # Add backward hooks to the parameters of the current layer if autograd is enabled.
      if torch.is_grad_enabled():
          self._add_backward_hooks()
  
      # Compute the outputs of the current layer
      res = self.module(*args, **kwargs)
  
      # Cleanup the parameters of the layer.
      #
      # *Skip cleaning up if autograd is enabled and this is the last layer in the network,
      # because we will need to fetch the parameters again for the backward pass.*
      if not torch.is_grad_enabled() or self.next_layer:
          self._cleanup_params()
  
      return res
  
imports: import functools
  from typing import List, Optional, Tuple
  import torch
  import torch.distributed
  from torch import nn
label_desc: Use image node_id 7 for calling the Zero3Layer obj's underlying member method code with example usage: obj.forward() and returns: res