code: def __init__(
      self,
      hidden_size=768,
      num_layers=12,
      num_attention_heads=12,
      max_seq_len=1024,
      vocab_size=50257,
      checkpoint=False,
  ):
      super().__init__()
      self.checkpoint = checkpoint
      self.config = GPT2Config(
          n_embd=hidden_size,
          n_layer=num_layers,
          n_head=num_attention_heads,
          n_positions=max_seq_len,
          n_ctx=max_seq_len,
          vocab_size=vocab_size,
      )
      self.model = GPT2LMHeadModel(self.config)
      if checkpoint:
          self.model.gradient_checkpointing_enable()
  
file_path: ColossalAI/examples/language/gpt/experiments/pipeline_parallel/model_zoo.py
imports: from torch import nn
  from transformers import GPT2Config, GPT2LMHeadModel
label_desc: Use image node_id 1 to create a new GPTLMModel object from inherited base classes: nn with example: obj = GPTLMModel(hidden_size, num_layers, num_attention_heads, max_seq_len, vocab_size, checkpoint)