src/transformers/models/jukebox/modeling_jukebox.py/JukeboxLayerStack/forward
class JukeboxLayerStack:    def forward(self, hidden_states, last_encoder_hidden_states=None, sample=False):
        # Blocks
        for i, attn_layer in enumerate(self._attn_mods):
            if attn_layer.attn_func == "cross_attention":  # attend to the lyrics
                hidden_states = attn_layer(
                    hidden_states, last_encoder_hidden_states=last_encoder_hidden_states, sample=sample
                )
            else:
                hidden_states = attn_layer(hidden_states, last_encoder_hidden_states=None, sample=sample)
            if attn_layer.attn.record_attn:
                self.saved_attn_weights.append(attn_layer.attn.c_attn.weight)
        return hidden_states