docstr: Return the working gradient for the specified parameter.
  
  Args:
      param_id (int): The index of the parameter.
  
  Returns:
      Tensor: The the working gradient slices for the specified param_id.
docstr_tok: ["Return","the","working","gradient","for","the","specified","parameter",".","Args",":","param_id","(","int",")",":","The","index","of","the","parameter",".","Returns",":","Tensor",":","The","the","working","gradient","slices","for","the","specified","param_id","."]
code: def get_working_grad_by_param_id(self, param_id) -> Tensor:
      """
      Return the working gradient for the specified parameter.
  
      Args:
          param_id (int): The index of the parameter.
  
      Returns:
          Tensor: The the working gradient slices for the specified param_id.
      """
  
      for group in self._grads_of_params.values():
          if param_id in group.keys():
              return group[param_id][self._working_index]
  
      raise KeyError(
          f"Working gradient for param_id {param_id} not found."
      )
  
file_path: ColossalAI/colossalai/zero/low_level/bookkeeping/gradient_store.py
imports: from typing import List
  from torch import Tensor
  from .base_store import BaseStore
label_desc: Use image node_id 6 for calling the GradientStore obj's underlying member method code with example usage: obj.get_working_grad_by_param_id(param_id) and returns: group