src/transformers/models/reformer/modeling_reformer.py/ReformerOnlyLMHead/forward_chunk
class ReformerOnlyLMHead:    def forward_chunk(self, hidden_states):
        hidden_states = self.decoder(hidden_states)
        return hidden_states