code: def _transform(
      self,
      transformer,
      global_ids: List[torch.Tensor],
      cache_ids: List[torch.Tensor],
  ):
      with torch.cuda.stream(self._stream):
          total_numel = sum([tensor.numel() for tensor in global_ids])
          if total_numel > 1e6:
              all_tensor = torch.cat(global_ids).to("cuda:0")
              unique_all_tensor, index = torch.unique(
                  all_tensor, return_inverse=True
              )
              unique_all_tensor = unique_all_tensor.to("cpu")
              all_cache = torch.empty_like(unique_all_tensor)
              success, ids_to_fetch = transformer.transform(
                  TensorList([unique_all_tensor]),
                  TensorList([all_cache]),
                  self._time,
              )
              del all_tensor
              all_tensor = torch.take(all_cache.to("cuda:0"), index)
              offset = 0
              for tensor in cache_ids:
                  numel = tensor.numel()
                  tensor.copy_(all_tensor[offset : offset + numel])
                  offset += numel
              assert (
                  total_numel == offset
              ), f"total_numel not equal offset, {total_numel} vs {offset}"
          else:
              # broadcast result
              success, ids_to_fetch = transformer.transform(
                  TensorList(global_ids),
                  TensorList(cache_ids),
                  self._time,
              )
          return success, ids_to_fetch
  
file_path: torchrec/contrib/dynamic_embedding/src/torchrec_dynamic_embedding/id_transformer_collection.py
imports: from typing import List, Tuple, Union
  import torch
  import torch.distributed
  from torchrec import EmbeddingBagConfig, EmbeddingConfig, KeyedJaggedTensor
  from .distributed import broadcast_ids_to_evict, broadcast_transform_result, gather_global_ids, scatter_cache_ids
  from .id_transformer import IDTransformer, TensorList
  from .ps import PSCollection
label_desc: Use image node_id 2 for calling the IDTransformerCollection obj's underlying member method code with example usage: obj._transform(transformer, global_ids, cache_ids) and returns: success,  ids_to_fetch