src/transformers/models/jukebox/modeling_jukebox.py/JukeboxVQVAE/sample
class JukeboxVQVAE:    def sample(self, n_samples):
        music_tokens = [
            torch.randint(0, self.nb_discrete_codes, size=(n_samples, *music_tokens_shape), device="cpu")
            for music_tokens_shape in self.music_tokens_shapes
        ]
        return self.decode(music_tokens)