docstr: Init grad chunk. This should be called in grad handler.
  
  Returns:
      Chunk: Grad chunk
docstr_tok: ["Init","grad","chunk",".","This","should","be","called","in","grad","handler",".","Returns",":","Chunk",":","Grad","chunk"]
code: def init_grad_chunk(self) -> "Chunk":
      """Init grad chunk. This should be called in grad handler.
  
      Returns:
          Chunk: Grad chunk
      """
      if self.grad_chunk is None:
          # grad chunk is not initialized
          grad_chunk = Chunk(
              chunk_size=self.chunk_size,
              zero_group=self.torch_pg,
              dtype=self.dtype,
              keep_gathered=self.keep_gathered,
              pin_memory=self.pin_memory,
              extra_dp_group=self.extra_dp_group,
          )
          grad_chunk.num_tensors = self.num_tensors
          grad_chunk.utilized_size = self.utilized_size
          grad_chunk.tensor_state_cnter[
              TensorState.HOLD
          ] = self.num_tensors
          for tensor, state in self.tensors_info.items():
              grad_chunk.tensors_info[tensor] = TensorInfo(
                  TensorState.HOLD, state.offset, state.end
              )
  
          grad_chunk.valid_end = self.valid_end
  
          if grad_chunk.chunk_temp.device.type == "cpu":
              grad_chunk.cuda_global_chunk = grad_chunk.chunk_temp.to(
                  get_current_device()
              )
          else:
              grad_chunk.cuda_global_chunk = grad_chunk.chunk_temp
          grad_chunk.chunk_temp = None
  
          if grad_chunk.pin_memory:
              grad_chunk.cpu_shard = torch.empty(
                  grad_chunk.shard_size,
                  dtype=grad_chunk.dtype,
                  pin_memory=grad_chunk.pin_memory,
              )
  
          self.grad_chunk = grad_chunk
      else:
          # grad chunk is initialized, just reallocate cuda global chunk
          self.grad_chunk.cuda_shard = None
          self.grad_chunk.is_gathered = True
          self.grad_chunk.l2_norm = None
          alloc_storage(self.grad_chunk.cuda_global_chunk)
  
      return self.grad_chunk
  
imports: from dataclasses import dataclass
  from enum import Enum
  from typing import Dict, List, Optional
  import torch
  import torch.distributed
  from torch.distributed import ProcessGroup
  from colossalai.utils import get_current_device
  from colossalai.utils.device import IS_NPU_AVAILABLE
label_desc: Use image node_id 33 for calling the Chunk obj's underlying member method code with example usage: obj.init_grad_chunk() and returns: self