src/transformers/models/jukebox/modeling_jukebox.py/JukeboxBlock/forward
class JukeboxBlock:    def forward(self, hidden_states, last_encoder_hidden_states, sample=False):
        residuals = hidden_states
        hidden_states = self.layer_norm_0(hidden_states)
        hidden_states = self.attn(hidden_states, last_encoder_hidden_states, sample)

        output_states = self.layer_norm_1(residuals + hidden_states)
        output_states = self.mlp(output_states)
        if self.res_scale == 1.0:
            output = residuals + hidden_states + output_states
        else:
            output = residuals + self.res_scale * (hidden_states + output_states)
        return output