src/transformers/models/reformer/modeling_reformer.py/ReformerOnlyLMHead/forward
class ReformerOnlyLMHead:    def forward(self, hidden_states):
        return apply_chunking_to_forward(self.forward_chunk, self.chunk_size_lm_head, self.seq_len_dim, hidden_states)