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

        n_states = int(np.ceil(sequence_length / self.transition_dim))

        offsets = torch.arange(self.transition_dim) * self.vocab_size
        offsets = offsets.repeat(n_states).to(trajectories.device)

        offset_trajectories = trajectories + offsets[:sequence_length]
        offset_trajectories[trajectories == self.vocab_size] = self.stop_token
        return offset_trajectories