def METHOD_NAME(
    self, tokens: torch.LongTensor, return_all_hiddens: bool = False
) -> torch.Tensor:
    if tokens.dim() == 1:
        tokens = tokens.unsqueeze(0)
    if tokens.size(-1) > min(self.model.max_positions()):
        raise ValueError(
            "tokens exceeds maximum length: {} > {}".format(
                tokens.size(-1), self.model.max_positions()
            )
        )
    tokens.to(device=self.device),
    prev_output_tokens = tokens.clone()
    prev_output_tokens[:, 0] = tokens.gather(
        1,
        (tokens.ne(self.task.source_dictionary.pad()).sum(dim=1) - 1).unsqueeze(-1),
    ).squeeze()
    prev_output_tokens[:, 1:] = tokens[:, :-1]
    features, extra = self.model(
        src_tokens=tokens,
        src_lengths=None,
        prev_output_tokens=prev_output_tokens,
        features_only=True,
        return_all_hiddens=return_all_hiddens,
    )
    if return_all_hiddens:
        # convert from T x B x C -> B x T x C
        inner_states = extra["inner_states"]
        return [inner_state.transpose(0, 1) for inner_state in inner_states]
    else:
        return features  # just the last layer's features