src/transformers/models/jukebox/modeling_jukebox.py/JukeboxBottleneckBlock/encode
class JukeboxBottleneckBlock:    def encode(self, latent_states):
        samples, _, seq_len = latent_states.shape

        # Preprocess.
        latent_states, _ = self.preprocess(latent_states)

        # Quantise
        music_tokens, _ = self.quantise(latent_states)

        # Postprocess.
        music_tokens = music_tokens.view(samples, seq_len)
        return music_tokens