src/transformers/models/trajectory_transformer/modeling_trajectory_transformer.py/TrajectoryTransformerModel/pad_to_full_observation
class TrajectoryTransformerModel:    def pad_to_full_observation(self, hidden_states):
        batch_size, sequence_length, _ = hidden_states.shape

        n_pad = (self.transition_dim - sequence_length % self.transition_dim) % self.transition_dim
        padding = torch.zeros(batch_size, n_pad, self.embedding_dim, device=hidden_states.device)

        # [ batch_size x padded_sequence_length' x embedding_dim ]
        hidden_states_pad = torch.cat([hidden_states, padding], dim=1)
        hidden_states_pad = hidden_states_pad.view(-1, self.transition_dim, self.embedding_dim)

        return hidden_states_pad, n_pad