code: def _weight_alloc(self, dtype, device):
      weight = torch.empty(
          self.num_embeddings,
          self.embedding_dim_per_partition,
          device=device,
          dtype=dtype,
      )
      with torch.no_grad():
          weight.data.uniform_(
              -1 / self.num_embeddings, 1 / self.num_embeddings
          )
          if self.padding_idx is not None:
              weight[self.padding_idx].fill_(0)
      colo_tensor_spec = ColoTensorSpec(
          pg=ProcessGroup(tp_degree=self.world_size),
          dist_attr=ShardSpec(
              dims=[-1], num_partitions=[self.world_size]
          ),
          compute_attr=ComputePattern.TP1D,
      )
      return ColoTensor.from_torch_tensor(weight, spec=colo_tensor_spec)
  
file_path: ColossalAI/colossalai/legacy/nn/parallel/layers/cache_embedding/parallel_cached_embedding.py
imports: from typing import List, Optional, Tuple
  import torch
  import torch.nn.functional
  from colossalai.legacy.nn._ops._utils import dual_all_to_all
  from colossalai.legacy.tensor import ColoTensorSpec, ComputePattern, ProcessGroup, ShardSpec
  from colossalai.tensor import ColoTensor
  from .cache_mgr import EvictionStrategy
  from .cached_embedding import CachedEmbeddingBag
label_desc: Use image node_id 2 for calling the ParallelCachedEmbeddingBag obj's underlying member method code with example usage: obj._weight_alloc(dtype, device) and returns: ColoTensor