class_docstr: the full GPT language model, with a context size of block_size
docstr: Returns:
  
  Examples:
  
  ```python
  >>> from transformers import TrajectoryTransformerModel
  >>> import torch
  
  >>> model = TrajectoryTransformerModel.from_pretrained(
  ...     "CarlCochet/trajectory-transformer-halfcheetah-medium-v2"
  ... )
  >>> model.to(device)
  >>> model.eval()
  
  >>> observations_dim, action_dim, batch_size = 17, 6, 256
  >>> seq_length = observations_dim + action_dim + 1
  
  >>> trajectories = torch.LongTensor([np.random.permutation(self.seq_length) for _ in range(batch_size)]).to(
  ...     device
  ... )
  >>> targets = torch.LongTensor([np.random.permutation(self.seq_length) for _ in range(batch_size)]).to(device)
  
  >>> outputs = model(
  ...     trajectories,
  ...     targets=targets,
  ...     use_cache=True,
  ...     output_attentions=True,
  ...     output_hidden_states=True,
  ...     return_dict=True,
  ... )
  ```
code: def forward(
      self,
      trajectories: Optional[torch.LongTensor] = None,
      past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
      targets: Optional[torch.FloatTensor] = None,
      attention_mask: Optional[torch.FloatTensor] = None,
      use_cache: Optional[bool] = None,
      output_attentions: Optional[bool] = None,
      output_hidden_states: Optional[bool] = None,
      return_dict: Optional[bool] = None,
  ) -> Union[Tuple[torch.Tensor], TrajectoryTransformerOutput]:
      r"""
      Returns:
  
      Examples:
  
      ```python
      >>> from transformers import TrajectoryTransformerModel
      >>> import torch
  
      >>> model = TrajectoryTransformerModel.from_pretrained(
      ...     "CarlCochet/trajectory-transformer-halfcheetah-medium-v2"
      ... )
      >>> model.to(device)
      >>> model.eval()
  
      >>> observations_dim, action_dim, batch_size = 17, 6, 256
      >>> seq_length = observations_dim + action_dim + 1
  
      >>> trajectories = torch.LongTensor([np.random.permutation(self.seq_length) for _ in range(batch_size)]).to(
      ...     device
      ... )
      >>> targets = torch.LongTensor([np.random.permutation(self.seq_length) for _ in range(batch_size)]).to(device)
  
      >>> outputs = model(
      ...     trajectories,
      ...     targets=targets,
      ...     use_cache=True,
      ...     output_attentions=True,
      ...     output_hidden_states=True,
      ...     return_dict=True,
      ... )
      ```
      """
      output_attentions = (
          output_attentions
          if output_attentions is not None
          else self.config.output_attentions
      )
      output_hidden_states = (
          output_hidden_states
          if output_hidden_states is not None
          else self.config.output_hidden_states
      )
  
      if past_key_values is None:
          past_key_values = tuple([None] * len(self.blocks))
  
      batch_size, sequence_length = trajectories.size()
  
      if sequence_length > self.block_size:
          raise ValueError(
              "Cannot forward, model block size is exhausted."
          )
  
      offset_trajectories = self.offset_tokens(trajectories)
      # [ batch_size x sequence_length x embedding_dim ]
      # forward the GPT model
      token_embeddings = self.tok_emb(
          offset_trajectories
      )  # each index maps to a (learnable) vector
      position_embeddings = self.pos_emb[
          :, :sequence_length, :
      ]  # each position maps to a (learnable) vector
  
      hidden_states = self.drop(token_embeddings + position_embeddings)
  
      if self.gradient_checkpointing and self.training:
          if use_cache:
              logger.warning_once(
                  "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
              )
              use_cache = False
  
      presents = () if use_cache else None
      all_self_attentions = () if output_attentions else None
      all_hidden_states = () if output_hidden_states else None
  
      for i, (block, layer_past) in enumerate(
          zip(self.blocks, past_key_values)
      ):
          if output_hidden_states:
              all_hidden_states = all_hidden_states + (hidden_states,)
  
          if self.gradient_checkpointing and self.training:
              outputs = self._gradient_checkpointing_func(
                  block.__call__,
                  hidden_states,
                  layer_past,
                  use_cache,
                  output_attentions,
              )
          else:
              outputs = block(
                  hidden_states,
                  layer_past,
                  use_cache,
                  output_attentions,
              )
  
          hidden_states = outputs[0]
          if use_cache is True:
              presents = presents + (outputs[1],)
  
          if output_attentions:
              all_self_attentions = all_self_attentions + (
                  outputs[2 if use_cache else 1],
              )
  
      # [ batch_size x sequence_length x embedding_dim ]
      hidden_state = self.ln_f(hidden_states)
  
      if output_hidden_states:
          all_hidden_states = all_hidden_states + (hidden_states,)
  
      hidden_states_pad, n_pad = self.pad_to_full_observation(
          hidden_state
      )
  
      logits = self.head(hidden_states_pad)
      logits = logits.reshape(
          batch_size, sequence_length + n_pad, self.vocab_size + 1
      )
      logits = logits[:, :sequence_length]
  
      # if we are given some desired targets also calculate the loss
      if targets is not None:
          loss = F.cross_entropy(
              logits.reshape(-1, logits.size(-1)),
              targets.view(-1),
              reduction="none",
          )
          if (
              self.action_weight != 1
              or self.reward_weight != 1
              or self.value_weight != 1
          ):
              # make weights
              n_states = int(
                  np.ceil(sequence_length / self.transition_dim)
              )
              weights = torch.cat(
                  [
                      torch.ones(
                          self.observation_dim,
                          device=trajectories.device,
                      ),
                      torch.ones(
                          self.action_dim, device=trajectories.device
                      )
                      * self.action_weight,
                      torch.ones(1, device=trajectories.device)
                      * self.reward_weight,
                      torch.ones(1, device=trajectories.device)
                      * self.value_weight,
                  ]
              )
              weights = weights.repeat(n_states)
              weights = weights[1:].repeat(batch_size, 1)
              loss = loss * weights.view(-1)
          loss = (loss * attention_mask.view(-1)).mean()
      else:
          loss = None
  
      if not return_dict:
          return tuple(
              v
              for v in [
                  loss,
                  logits,
                  presents,
                  all_hidden_states,
                  all_self_attentions,
              ]
              if v is not None
          )
  
      return TrajectoryTransformerOutput(
          loss=loss,
          logits=logits,
          past_key_values=presents,
          hidden_states=all_hidden_states,
          attentions=all_self_attentions,
      )
  
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 5 for calling the TrajectoryTransformerModel obj's underlying member method code with example usage: obj.forward(trajectories, past_key_values, targets, attention_mask, use_cache, output_attentions, output_hidden_states, return_dict) and returns: TrajectoryTransformerOutput, tuple