code: def data_gen(batch_size: int = 4, seq_len: int = 512):
      input_ids = torch.randint(
          10, 30000, (1, seq_len), dtype=torch.int32
      )
      attention_mask = torch.ones((1, seq_len), dtype=torch.int32)
      data = dict(input_ids=input_ids, attention_mask=attention_mask)
      for k, v in data.items():
          if torch.is_tensor(v) or "Tensor" in v.__class__.__name__:
              new_shape = [1] * v.dim()
              new_shape[0] = batch_size
              data[k] = v.to("cuda").repeat(*new_shape)
      return data
  
file_path: ColossalAI/colossalai/legacy/inference/pipeline/benchmark/benchmark.py
imports: import argparse
  import time
  import torch
  import torch.distributed
  import transformers
  import colossalai
  from colossalai.inference import PPInferEngine
  from colossalai.inference.pipeline.policies import LlamaModelInferPolicy
label_desc: Use image node_id 1 for calling a global function with example usage: data_gen(batch_size, seq_len) and returns: data