code: def data_gen():
      return dict(x=torch.randint(low=0, high=20, size=(16,)))
  
file_path: ColossalAI/tests/kit/model_zoo/custom/simple_net.py
imports: import torch
  import torch.nn
  import torch.nn.functional
  from ..registry import model_zoo
  from .base import CheckpointModule
label_desc: Use image node_id 1 for calling a global function with example usage: data_gen() and returns: dict