class_docstr: the full GPT language model, with a context size of block_size
code: def __init__(self, config):
      super().__init__(config)
  
      # input embedding stem (+1 for stop token)
      self.tok_emb = nn.Embedding(
          config.vocab_size * config.transition_dim + 1, config.n_embd
      )
  
      self.pos_emb = nn.Parameter(
          torch.zeros(1, config.block_size, config.n_embd)
      )
      self.drop = nn.Dropout(config.embd_pdrop)
      # transformer
      self.blocks = nn.ModuleList(
          [Block(config) for _ in range(config.n_layer)]
      )
      # decoder head
      self.ln_f = nn.LayerNorm(config.n_embd)
      self.head = EinLinear(
          config.transition_dim,
          config.n_embd,
          config.vocab_size + 1,
          bias=False,
      )
  
      self.vocab_size = config.vocab_size
      self.stop_token = config.vocab_size * config.transition_dim
      self.block_size = config.block_size
  
      self.observation_dim = config.observation_dim
      self.action_dim = config.action_dim
      self.transition_dim = config.transition_dim
      self.embedding_dim = config.n_embd
  
      self.action_weight = config.action_weight
      self.reward_weight = config.reward_weight
      self.value_weight = config.value_weight
  
      self.gradient_checkpointing = False
  
      self.post_init()
  
imports: import math
  import os
  from dataclasses import dataclass
  from typing import Optional, Tuple, Union
  import numpy
  import torch
  import torch.utils.checkpoint
  from torch import nn
  from torch.nn import functional
  from ....modeling_utils import PreTrainedModel
  from ....utils import ModelOutput, add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
  from .configuration_trajectory_transformer import TrajectoryTransformerConfig
label_desc: Use image node_id 1 to create a new TrajectoryTransformerModel object from inherited base classes: TrajectoryTransformerPreTrainedModel with example: obj = TrajectoryTransformerModel(config)