code: def get_data(batch_size, seq_len, vocab_size):
      input_ids = torch.randint(
          0,
          vocab_size,
          (batch_size, seq_len),
          device=torch.cuda.current_device(),
      )
      attention_mask = torch.ones_like(input_ids)
      return input_ids, attention_mask
  
file_path: ColossalAI/examples/language/gpt/experiments/auto_parallel/auto_parallel_with_gpt.py
imports: from functools import partial
  from time import time
  import psutil
  import torch
  import transformers
  from gpt_modules import GPT2LMHeadModel, GPTLMLoss
  from colossalai.auto_parallel.tensor_shard.initialize import autoparallelize
  from colossalai.initialize import launch_from_torch
  from colossalai.legacy.core import global_context
  from colossalai.logging import disable_existing_loggers, get_dist_logger
label_desc: Use image node_id 5 for calling a global function with example usage: get_data(batch_size, seq_len, vocab_size) and returns: input_ids,  attention_mask