code: def output_transform_fn(x):
      if isinstance(x, KeyedTensor):
          output = dict()
          for key in x.keys():
              output[key] = x[key]
          return output
      else:
          return dict(output=x)
  
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 6 for calling a global function with example usage: output_transform_fn(x) and returns: output, dict