code: def data_gen():
      # Generated from following code snippet
      #
      # from transformers import BertTokenizer
      # input = 'Hello, my dog is cute'
      # tokenized_input = tokenizer(input, return_tensors='pt')
      # input_ids = tokenized_input['input_ids']
      # attention_mask = tokenized_input['attention_mask']
      # token_type_ids = tokenized_input['token_type_ids']
      input_ids = torch.tensor(
          [[101, 7592, 1010, 2026, 3899, 2003, 10140, 102]],
          dtype=torch.int64,
      )
      token_type_ids = torch.tensor(
          [[0, 0, 0, 0, 0, 0, 0, 0]], dtype=torch.int64
      )
      attention_mask = torch.tensor(
          [[1, 1, 1, 1, 1, 1, 1, 0]], dtype=torch.int64
      )
      return dict(
          input_ids=input_ids,
          token_type_ids=token_type_ids,
          attention_mask=attention_mask,
      )
  
file_path: colossalai/tests/kit/model_zoo/transformers/bert.py
imports: import torch
  import transformers
  from ..registry import ModelAttribute, model_zoo
label_desc: Use image node_id 1 for calling a global function with example usage: data_gen() and returns: dict