code: def forward(
      self,
      indices,
      offsets=None,
      per_sample_weights=None,
      shape_hook=None,
      scatter_dim=0,
      gather_dim=-1,
  ):
      if self.cache_op:
          with torch.no_grad():
              indices = self.cache_weight_mgr.prepare_ids(indices)
      output_shard = F.embedding_bag(
          indices.cuda(),
          self.cache_weight_mgr.cuda_cached_weight,
          offsets,
          self.max_norm,
          self.norm_type,
          self.scale_grad_by_freq,
          self.mode,
          self.sparse,
          per_sample_weights,
          self.include_last_offset,
          self.padding_idx,
      )
      if shape_hook is not None:
          output_shard = shape_hook(output_shard)
      output_full = dual_all_to_all(
          output_shard,
          self.weight.get_process_group(),
          scatter_dim=scatter_dim,
          gather_dim=gather_dim,
      )
      return output_full
  
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 3 for calling the ParallelCachedEmbeddingBag obj's underlying member method code with example usage: obj.forward(indices, offsets, per_sample_weights, shape_hook, scatter_dim, gather_dim) and returns: output_full