src/transformers/models/jukebox/modeling_jukebox.py/JukeboxMusicTokenConditioner/__init__
class JukeboxMusicTokenConditioner:    def __init__(self, config, level):
        super().__init__()
        self.embed_tokens = nn.Embedding(config.music_vocab_size, config.hidden_size)
        config.embed_dim = config.music_vocab_size  # setting correct argument for the `JukeboxDecoder`

        self.upsampler = JukeboxDecoderConvBock(
            config,
            config.hidden_size,
            config.res_conv_width,
            config.res_conv_depth,
            config.res_downs_t[level],
            config.res_strides_t[level],
            reverse_dilation=False,
        )
        self.layer_norm = JukeboxLayerNorm(config.hidden_size)