src/transformers/models/jukebox/modeling_jukebox.py/JukeboxEncoder/forward
class JukeboxEncoder:    def forward(self, hidden_states):
        all_hidden_states = []

        # 64, 32, ...
        for level in range(self.levels):
            level_block = self.level_blocks[level]
            hidden_states = level_block(hidden_states)
            all_hidden_states.append(hidden_states)

        return all_hidden_states