code: def get_temp_total_chunk_on_cuda(chunk: Chunk, dtype: torch.dtype):
      if chunk.is_gathered:
          return chunk.cuda_global_chunk
  
      if chunk.cuda_shard is not None:
          shard_temp = chunk.cuda_shard
      else:
          shard_temp = chunk.cpu_shard.to(get_current_device())
  
      shard_temp = shard_temp.to(dtype)
  
      total_temp = torch.zeros(
          chunk.chunk_size, dtype=dtype, device=get_current_device()
      )
      gather_list = list(
          torch.chunk(input=total_temp, chunks=chunk.pg_size, dim=0)
      )
      dist.all_gather(
          tensor_list=gather_list,
          tensor=shard_temp,
          group=chunk.torch_pg,
      )
  
      return total_temp
  
file_path: colossalai/colossalai/zero/gemini/utils.py
imports: from collections import OrderedDict
  from copy import copy
  from typing import Optional, Set
  import torch
  import torch.distributed
  import torch.nn
  from colossalai.utils import get_current_device
  from .chunk import Chunk
label_desc: Use image node_id 1 for calling a global function with example usage: get_temp_total_chunk_on_cuda(chunk, dtype) and returns: total_temp, chunk