src/transformers/models/jukebox/modeling_jukebox.py/JukeboxBottleneckBlock/postprocess
class JukeboxBottleneckBlock:    def postprocess(self, latent_states, dequantised_states, x_shape):
        batch_size, time = x_shape
        dequantised_states = dequantised_states.view(batch_size, time, -1).permute(0, 2, 1).contiguous()
        latent_states = latent_states.view(batch_size, time)
        return latent_states, dequantised_states