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: #### Cleanup the parameter data
  
  This will release all the memory used by the layer parameters.
docstr_tok: ["#","#","#","#","Cleanup","the","parameter","data","This","will","release","all","the","memory","used","by","the","layer","parameters","."]
code: def _cleanup_params(self):
      """
      #### Cleanup the parameter data
  
      This will release all the memory used by the layer parameters.
      """
  
      # Set the flag to indicate that the parameters are not fetched
      self.is_fetched = False
  
      # Iterate through all parameters
      for ps in self.param_refs:
          for p in ps:
              # Wait for operations on the parameters to complete before any new operations
              p.data.record_stream(torch.cuda.current_stream())
              # Check to make sure the parameter is not sharing storage with anything else
              assert (
                  p.data.storage_offset() == 0
              ), "The tensor is not the sole occupant of the storage."
              # Resize the storage to $0$. This will release the memory used by the parameter.
              #
              # **Setting `p.data` will not release the memory, since the autograd graph keeps a reference to it.**
              p.data.storage().resize_(
                  0
              )  # This is what actually clears the memory
              # Make sure the parameter has no gradient data
              assert p.grad is None, "Gradients should be None"
  
imports: import functools
  from typing import List, Optional, Tuple
  import torch
  import torch.distributed
  from torch import nn
label_desc: Use image node_id 5 for calling the Zero3Layer obj's underlying member method code with example usage: obj._cleanup_params() without return types