code: def gen_kjt():
      KJT = KeyedJaggedTensor.from_offsets_sync(
          keys=["f1", "f2"],
          values=torch.tensor([1, 2, 3, 4, 5, 6, 7, 8]),
          offsets=torch.tensor([0, 2, 4, 6, 8]),
      )
      return KJT
  
file_path: colossalai/tests/kit/model_zoo/torchrec/torchrec.py
imports: from functools import partial
  import torch
  from torchrec.models import deepfm, dlrm
  from torchrec.modules.embedding_configs import EmbeddingBagConfig
  from torchrec.modules.embedding_modules import EmbeddingBagCollection
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor, KeyedTensor
  from ..registry import model_zoo
label_desc: Use image node_id 2 for calling a global function with example usage: gen_kjt() and returns: KJT