RuntimeError: Error(s) in loading state_dict for SimplePrior

[marypilataki]

Hello,

Thank you for making your work public. I am having an issue when trying to extract a jukebox representation using the model "5b". My python script is identical to your main under /representations/jukebox where I am using a different dataset.

Please see the exact error below

0: Loading vqvae in eval mode
Loading artist IDs from /data/home/acw512/musicnet_vgg_multitask/lib/python3.8/site-packages/jukebox/data/ids/v2_artist_ids.txt
Loading artist IDs from /data/home/acw512/musicnet_vgg_multitask/lib/python3.8/site-packages/jukebox/data/ids/v2_genre_ids.txt
Level:2, Cond downsample:None, Raw to tokens:128, Sample length:1048576
0: Converting to fp16 params
Downloading from azure
Running  wget -O /data/home/acw512/.cache/jukebox/models/5b/prior_level_2.pth.tar https://openaipublic.azureedge.net/jukebox/models/5b/prior_level_2.pth.tar
Restored from /data/home/acw512/.cache/jukebox/models/5b/prior_level_2.pth.tar
Traceback (most recent call last):
  File "test_representation.py", line 139, in <module>
    top_prior = make_prior(hparams, vqvae, device)
  File "/data/home/acw512/musicnet_vgg_multitask/lib/python3.8/site-packages/jukebox/make_models.py", line 179, in make_prior
    restore_model(hps, prior, hps.restore_prior)
  File "/data/home/acw512/musicnet_vgg_multitask/lib/python3.8/site-packages/jukebox/make_models.py", line 61, in restore_model
    model.load_state_dict(checkpoint['model'])
  File "/data/home/acw512/musicnet_vgg_multitask/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1406, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for SimplePrior:
    Unexpected key(s) in state_dict: "prior.transformer._attn_mods.36.attn.c_attn.w", "prior.transformer._attn_mods.36.attn.c_attn.b", "prior.transformer._attn_mods.36.attn.c_proj.w", "prior.transformer._attn_mods.36.attn.c_proj.b", "prior.transformer._attn_mods.36.ln_0.weight", "prior.transformer._attn_mods.36.ln_0.bias", "prior.transformer._attn_mods.36.mlp.c_fc.w", "prior.transformer._attn_mods.36.mlp.c_fc.b", "prior.transformer._attn_mods.36.mlp.c_proj.w", "prior.transformer._attn_mods.36.mlp.c_proj.b", "prior.transformer._attn_mods.36.ln_1.weight", "prior.transformer._attn_mods.36.ln_1.bias", "prior.transformer._attn_mods.37.attn.c_attn.w", "prior.transformer._attn_mods.37.attn.c_attn.b", "prior.transformer._attn_mods.37.attn.c_proj.w", "prior.transformer._attn_mods.37.attn.c_proj.b", "prior.transformer._attn_mods.37.ln_0.weight", "prior.transformer._attn_mods.37.ln_0.bias", "prior.transformer._attn_mods.37.mlp.c_fc.w", "prior.transformer._attn_mods.37.mlp.c_fc.b", "prior.transformer._attn_mods.37.mlp.c_proj.w", "prior.transformer._attn_mods.37.mlp.c_proj.b", "prior.transformer._attn_mods.37.ln_1.weight", "prior.transformer._attn_mods.37.ln_1.bias", "prior.transformer._attn_mods.38.attn.c_attn.w", "prior.transformer._attn_mods.38.attn.c_attn.b", "prior.transformer._attn_mods.38.attn.c_proj.w", "prior.transformer._attn_mods.38.attn.c_proj.b", "prior.transformer._attn_mods.38.ln_0.weight", "prior.transformer._attn_mods.38.ln_0.bias", "prior.transformer._attn_mods.38.mlp.c_fc.w", "prior.transformer._attn_mods.38.mlp.c_fc.b", "prior.transformer._attn_mods.38.mlp.c_proj.w", "prior.transformer._attn_mods.38.mlp.c_proj.b", "prior.transformer._attn_mods.38.ln_1.weight", "prior.transformer._attn_mods.38.ln_1.bias", "prior.transformer._attn_mods.39.attn.c_attn.w", "prior.transformer._attn_mods.39.attn.c_attn.b", "prior.transformer._attn_mods.39.attn.c_proj.w", "prior.transformer._attn_mods.39.attn.c_proj.b", "prior.transformer._attn_mods.39.ln_0.weight", "prior.transformer._attn_mods.39.ln_0.bias", "prior.transformer._attn_mods.39.mlp.c_fc.w", "prior.transformer._attn_mods.39.mlp.c_fc.b", "prior.transformer._attn_mods.39.mlp.c_proj.w", "prior.transformer._attn_mods.39.mlp.c_proj.b", "prior.transformer._attn_mods.39.ln_1.weight", "prior.transformer._attn_mods.39.ln_1.bias", "prior.transformer._attn_mods.40.attn.c_attn.w", "prior.transformer._attn_mods.40.attn.c_attn.b", "prior.transformer._attn_mods.40.attn.c_proj.w", "prior.transformer._attn_mods.40.attn.c_proj.b", "prior.transformer._attn_mods.40.ln_0.weight", "prior.transformer._attn_mods.40.ln_0.bias", "prior.transformer._attn_mods.40.mlp.c_fc.w", "prior.transformer._attn_mods.40.mlp.c_fc.b", "prior.transformer._attn_mods.40.mlp.c_proj.w", "prior.transformer._attn_mods.40.mlp.c_proj.b", "prior.transformer._attn_mods.40.ln_1.weight", "prior.transformer._attn_mods.40.ln_1.bias", "prior.transformer._attn_mods.41.attn.c_attn.w", "prior.transformer._attn_mods.41.attn.c_attn.b", "prior.transformer._attn_mods.41.attn.c_proj.w", "prior.transformer._attn_mods.41.attn.c_proj.b", "prior.transformer._attn_mods.41.ln_0.weight", "prior.transformer._attn_mods.41.ln_0.bias", "prior.transformer._attn_mods.41.mlp.c_fc.w", "prior.transformer._attn_mods.41.mlp.c_fc.b", "prior.transformer._attn_mods.41.mlp.c_proj.w", "prior.transformer._attn_mods.41.mlp.c_proj.b", "prior.transformer._attn_mods.41.ln_1.weight", "prior.transformer._attn_mods.41.ln_1.bias", "prior.transformer._attn_mods.42.attn.c_attn.w", "prior.transformer._attn_mods.42.attn.c_attn.b", "prior.transformer._attn_mods.42.attn.c_proj.w", "prior.transformer._attn_mods.42.attn.c_proj.b", "prior.transformer._attn_mods.42.ln_0.weight", "prior.transformer._attn_mods.42.ln_0.bias", "prior.transformer._attn_mods.42.mlp.c_fc.w", "prior.transformer._attn_mods.42.mlp.c_fc.b", "prior.transformer._attn_mods.42.mlp.c_proj.w", "prior.transformer._attn_mods.42.mlp.c_proj.b", "prior.transformer._attn_mods.42.ln_1.weight", "prior.transformer._attn_mods.42.ln_1.bias", "prior.transformer._attn_mods.43.attn.c_attn.w", "prior.transformer._attn_mods.43.attn.c_attn.b", "prior.transformer._attn_mods.43.attn.c_proj.w", "prior.transformer._attn_mods.43.attn.c_proj.b", "prior.transformer._attn_mods.43.ln_0.weight", "prior.transformer._attn_mods.43.ln_0.bias", "prior.transformer._attn_mods.43.mlp.c_fc.w", "prior.transformer._attn_mods.43.mlp.c_fc.b", "prior.transformer._attn_mods.43.mlp.c_proj.w", "prior.transformer._attn_mods.43.mlp.c_proj.b", "prior.transformer._attn_mods.43.ln_1.weight", "prior.transformer._attn_mods.43.ln_1.bias", "prior.transformer._attn_mods.44.attn.c_attn.w", "prior.transformer._attn_mods.44.attn.c_attn.b", "prior.transformer._attn_mods.44.attn.c_proj.w", "prior.transformer._attn_mods.44.attn.c_proj.b", "prior.transformer._attn_mods.44.ln_0.weight", "prior.transformer._attn_mods.44.ln_0.bias", "prior.transformer._attn_mods.44.mlp.c_fc.w", "prior.transformer._attn_mods.44.mlp.c_fc.b", "prior.transformer._attn_mods.44.mlp.c_proj.w", "prior.transformer._attn_mods.44.mlp.c_proj.b", "prior.transformer._attn_mods.44.ln_1.weight", "prior.transformer._attn_mods.44.ln_1.bias", "prior.transformer._attn_mods.45.attn.c_attn.w", "prior.transformer._attn_mods.45.attn.c_attn.b", "prior.transformer._attn_mods.45.attn.c_proj.w", "prior.transformer._attn_mods.45.attn.c_proj.b", "prior.transformer._attn_mods.45.ln_0.weight", "prior.transformer._attn_mods.45.ln_0.bias", "prior.transformer._attn_mods.45.mlp.c_fc.w", "prior.transformer._attn_mods.45.mlp.c_fc.b", "prior.transformer._attn_mods.45.mlp.c_proj.w", "prior.transformer._attn_mods.45.mlp.c_proj.b", "prior.transformer._attn_mods.45.ln_1.weight", "prior.transformer._attn_mods.45.ln_1.bias", "prior.transformer._attn_mods.46.attn.c_attn.w", "prior.transformer._attn_mods.46.attn.c_attn.b", "prior.transformer._attn_mods.46.attn.c_proj.w", "prior.transformer._attn_mods.46.attn.c_proj.b", "prior.transformer._attn_mods.46.ln_0.weight", "prior.transformer._attn_mods.46.ln_0.bias", "prior.transformer._attn_mods.46.mlp.c_fc.w", "prior.transformer._attn_mods.46.mlp.c_fc.b", "prior.transformer._attn_mods.46.mlp.c_proj.w", "prior.transformer._attn_mods.46.mlp.c_proj.b", "prior.transformer._attn_mods.46.ln_1.weight", "prior.transformer._attn_mods.46.ln_1.bias", "prior.transformer._attn_mods.47.attn.c_attn.w", "prior.transformer._attn_mods.47.attn.c_attn.b", "prior.transformer._attn_mods.47.attn.c_proj.w", "prior.transformer._attn_mods.47.attn.c_proj.b", "prior.transformer._attn_mods.47.ln_0.weight", "prior.transformer._attn_mods.47.ln_0.bias", "prior.transformer._attn_mods.47.mlp.c_fc.w", "prior.transformer._attn_mods.47.mlp.c_fc.b", "prior.transformer._attn_mods.47.mlp.c_proj.w", "prior.transformer._attn_mods.47.mlp.c_proj.b", "prior.transformer._attn_mods.47.ln_1.weight", "prior.transformer._attn_mods.47.ln_1.bias", "prior.transformer._attn_mods.48.attn.c_attn.w", "prior.transformer._attn_mods.48.attn.c_attn.b", "prior.transformer._attn_mods.48.attn.c_proj.w", "prior.transformer._attn_mods.48.attn.c_proj.b", "prior.transformer._attn_mods.48.ln_0.weight", "prior.transformer._attn_mods.48.ln_0.bias", "prior.transformer._attn_mods.48.mlp.c_fc.w", "prior.transformer._attn_mods.48.mlp.c_fc.b", "prior.transformer._attn_mods.48.mlp.c_proj.w", "prior.transformer._attn_mods.48.mlp.c_proj.b", "prior.transformer._attn_mods.48.ln_1.weight", "prior.transformer._attn_mods.48.ln_1.bias", "prior.transformer._attn_mods.49.attn.c_attn.w", "prior.transformer._attn_mods.49.attn.c_attn.b", "prior.transformer._attn_mods.49.attn.c_proj.w", "prior.transformer._attn_mods.49.attn.c_proj.b", "prior.transformer._attn_mods.49.ln_0.weight", "prior.transformer._attn_mods.49.ln_0.bias", "prior.transformer._attn_mods.49.mlp.c_fc.w", "prior.transformer._attn_mods.49.mlp.c_fc.b", "prior.transformer._attn_mods.49.mlp.c_proj.w", "prior.transformer._attn_mods.49.mlp.c_proj.b", "prior.transformer._attn_mods.49.ln_1.weight", "prior.transformer._attn_mods.49.ln_1.bias", "prior.transformer._attn_mods.50.attn.c_attn.w", "prior.transformer._attn_mods.50.attn.c_attn.b", "prior.transformer._attn_mods.50.attn.c_proj.w", "prior.transformer._attn_mods.50.attn.c_proj.b", "prior.transformer._attn_mods.50.ln_0.weight", "prior.transformer._attn_mods.50.ln_0.bias", "prior.transformer._attn_mods.50.mlp.c_fc.w", "prior.transformer._attn_mods.50.mlp.c_fc.b", "prior.transformer._attn_mods.50.mlp.c_proj.w", "prior.transformer._attn_mods.50.mlp.c_proj.b", "prior.transformer._attn_mods.50.ln_1.weight", "prior.transformer._attn_mods.50.ln_1.bias", "prior.transformer._attn_mods.51.attn.c_attn.w", "prior.transformer._attn_mods.51.attn.c_attn.b", "prior.transformer._attn_mods.51.attn.c_proj.w", "prior.transformer._attn_mods.51.attn.c_proj.b", "prior.transformer._attn_mods.51.ln_0.weight", "prior.transformer._attn_mods.51.ln_0.bias", "prior.transformer._attn_mods.51.mlp.c_fc.w", "prior.transformer._attn_mods.51.mlp.c_fc.b", "prior.transformer._attn_mods.51.mlp.c_proj.w", "prior.transformer._attn_mods.51.mlp.c_proj.b", "prior.transformer._attn_mods.51.ln_1.weight", "prior.transformer._attn_mods.51.ln_1.bias", "prior.transformer._attn_mods.52.attn.c_attn.w", "prior.transformer._attn_mods.52.attn.c_attn.b", "prior.transformer._attn_mods.52.attn.c_proj.w", "prior.transformer._attn_mods.52.attn.c_proj.b", "prior.transformer._attn_mods.52.ln_0.weight", "prior.transformer._attn_mods.52.ln_0.bias", "prior.transformer._attn_mods.52.mlp.c_fc.w", "prior.transformer._attn_mods.52.mlp.c_fc.b", "prior.transformer._attn_mods.52.mlp.c_proj.w", "prior.transformer._attn_mods.52.mlp.c_proj.b", "prior.transformer._attn_mods.52.ln_1.weight", "prior.transformer._attn_mods.52.ln_1.bias", "prior.transformer._attn_mods.53.attn.c_attn.w", "prior.transformer._attn_mods.53.attn.c_attn.b", "prior.transformer._attn_mods.53.attn.c_proj.w", "prior.transformer._attn_mods.53.attn.c_proj.b", "prior.transformer._attn_mods.53.ln_0.weight", "prior.transformer._attn_mods.53.ln_0.bias", "prior.transformer._attn_mods.53.mlp.c_fc.w", "prior.transformer._attn_mods.53.mlp.c_fc.b", "prior.transformer._attn_mods.53.mlp.c_proj.w", "prior.transformer._attn_mods.53.mlp.c_proj.b", "prior.transformer._attn_mods.53.ln_1.weight", "prior.transformer._attn_mods.53.ln_1.bias", "prior.transformer._attn_mods.54.attn.c_attn.w", "prior.transformer._attn_mods.54.attn.c_attn.b", "prior.transformer._attn_mods.54.attn.c_proj.w", "prior.transformer._attn_mods.54.attn.c_proj.b", "prior.transformer._attn_mods.54.ln_0.weight", "prior.transformer._attn_mods.54.ln_0.bias", "prior.transformer._attn_mods.54.mlp.c_fc.w", "prior.transformer._attn_mods.54.mlp.c_fc.b", "prior.transformer._attn_mods.54.mlp.c_proj.w", "prior.transformer._attn_mods.54.mlp.c_proj.b", "prior.transformer._attn_mods.54.ln_1.weight", "prior.transformer._attn_mods.54.ln_1.bias", "prior.transformer._attn_mods.55.attn.c_attn.w", "prior.transformer._attn_mods.55.attn.c_attn.b", "prior.transformer._attn_mods.55.attn.c_proj.w", "prior.transformer._attn_mods.55.attn.c_proj.b", "prior.transformer._attn_mods.55.ln_0.weight", "prior.transformer._attn_mods.55.ln_0.bias", "prior.transformer._attn_mods.55.mlp.c_fc.w", "prior.transformer._attn_mods.55.mlp.c_fc.b", "prior.transformer._attn_mods.55.mlp.c_proj.w", "prior.transformer._attn_mods.55.mlp.c_proj.b", "prior.transformer._attn_mods.55.ln_1.weight", "prior.transformer._attn_mods.55.ln_1.bias", "prior.transformer._attn_mods.56.attn.c_attn.w", "prior.transformer._attn_mods.56.attn.c_attn.b", "prior.transformer._attn_mods.56.attn.c_proj.w", "prior.transformer._attn_mods.56.attn.c_proj.b", "prior.transformer._attn_mods.56.ln_0.weight", "prior.transformer._attn_mods.56.ln_0.bias", "prior.transformer._attn_mods.56.mlp.c_fc.w", "prior.transformer._attn_mods.56.mlp.c_fc.b", "prior.transformer._attn_mods.56.mlp.c_proj.w", "prior.transformer._attn_mods.56.mlp.c_proj.b", "prior.transformer._attn_mods.56.ln_1.weight", "prior.transformer._attn_mods.56.ln_1.bias", "prior.transformer._attn_mods.57.attn.c_attn.w", "prior.transformer._attn_mods.57.attn.c_attn.b", "prior.transformer._attn_mods.57.attn.c_proj.w", "prior.transformer._attn_mods.57.attn.c_proj.b", "prior.transformer._attn_mods.57.ln_0.weight", "prior.transformer._attn_mods.57.ln_0.bias", "prior.transformer._attn_mods.57.mlp.c_fc.w", "prior.transformer._attn_mods.57.mlp.c_fc.b", "prior.transformer._attn_mods.57.mlp.c_proj.w", "prior.transformer._attn_mods.57.mlp.c_proj.b", "prior.transformer._attn_mods.57.ln_1.weight", "prior.transformer._attn_mods.57.ln_1.bias", "prior.transformer._attn_mods.58.attn.c_attn.w", "prior.transformer._attn_mods.58.attn.c_attn.b", "prior.transformer._attn_mods.58.attn.c_proj.w", "prior.transformer._attn_mods.58.attn.c_proj.b", "prior.transformer._attn_mods.58.ln_0.weight", "prior.transformer._attn_mods.58.ln_0.bias", "prior.transformer._attn_mods.58.mlp.c_fc.w", "prior.transformer._attn_mods.58.mlp.c_fc.b", "prior.transformer._attn_mods.58.mlp.c_proj.w", "prior.transformer._attn_mods.58.mlp.c_proj.b", "prior.transformer._attn_mods.58.ln_1.weight", "prior.transformer._attn_mods.58.ln_1.bias", "prior.transformer._attn_mods.59.attn.c_attn.w", "prior.transformer._attn_mods.59.attn.c_attn.b", "prior.transformer._attn_mods.59.attn.c_proj.w", "prior.transformer._attn_mods.59.attn.c_proj.b", "prior.transformer._attn_mods.59.ln_0.weight", "prior.transformer._attn_mods.59.ln_0.bias", "prior.transformer._attn_mods.59.mlp.c_fc.w", "prior.transformer._attn_mods.59.mlp.c_fc.b", "prior.transformer._attn_mods.59.mlp.c_proj.w", "prior.transformer._attn_mods.59.mlp.c_proj.b", "prior.transformer._attn_mods.59.ln_1.weight", "prior.transformer._attn_mods.59.ln_1.bias", "prior.transformer._attn_mods.60.attn.c_attn.w", "prior.transformer._attn_mods.60.attn.c_attn.b", "prior.transformer._attn_mods.60.attn.c_proj.w", "prior.transformer._attn_mods.60.attn.c_proj.b", "prior.transformer._attn_mods.60.ln_0.weight", "prior.transformer._attn_mods.60.ln_0.bias", "prior.transformer._attn_mods.60.mlp.c_fc.w", "prior.transformer._attn_mods.60.mlp.c_fc.b", "prior.transformer._attn_mods.60.mlp.c_proj.w", "prior.transformer._attn_mods.60.mlp.c_proj.b", "prior.transformer._attn_mods.60.ln_1.weight", "prior.transformer._attn_mods.60.ln_1.bias", "prior.transformer._attn_mods.61.attn.c_attn.w", "prior.transformer._attn_mods.61.attn.c_attn.b", "prior.transformer._attn_mods.61.attn.c_proj.w", "prior.transformer._attn_mods.61.attn.c_proj.b", "prior.transformer._attn_mods.61.ln_0.weight", "prior.transformer._attn_mods.61.ln_0.bias", "prior.transformer._attn_mods.61.mlp.c_fc.w", "prior.transformer._attn_mods.61.mlp.c_fc.b", "prior.transformer._attn_mods.61.mlp.c_proj.w", "prior.transformer._attn_mods.61.mlp.c_proj.b", "prior.transformer._attn_mods.61.ln_1.weight", "prior.transformer._attn_mods.61.ln_1.bias", "prior.transformer._attn_mods.62.attn.c_attn.w", "prior.transformer._attn_mods.62.attn.c_attn.b", "prior.transformer._attn_mods.62.attn.c_proj.w", "prior.transformer._attn_mods.62.attn.c_proj.b", "prior.transformer._attn_mods.62.ln_0.weight", "prior.transformer._attn_mods.62.ln_0.bias", "prior.transformer._attn_mods.62.mlp.c_fc.w", "prior.transformer._attn_mods.62.mlp.c_fc.b", "prior.transformer._attn_mods.62.mlp.c_proj.w", "prior.transformer._attn_mods.62.mlp.c_proj.b", "prior.transformer._attn_mods.62.ln_1.weight", "prior.transformer._attn_mods.62.ln_1.bias", "prior.transformer._attn_mods.63.attn.c_attn.w", "prior.transformer._attn_mods.63.attn.c_attn.b", "prior.transformer._attn_mods.63.attn.c_proj.w", "prior.transformer._attn_mods.63.attn.c_proj.b", "prior.transformer._attn_mods.63.ln_0.weight", "prior.transformer._attn_mods.63.ln_0.bias", "prior.transformer._attn_mods.63.mlp.c_fc.w", "prior.transformer._attn_mods.63.mlp.c_fc.b", "prior.transformer._attn_mods.63.mlp.c_proj.w", "prior.transformer._attn_mods.63.mlp.c_proj.b", "prior.transformer._attn_mods.63.ln_1.weight", "prior.transformer._attn_mods.63.ln_1.bias", "prior.transformer._attn_mods.64.attn.c_attn.w", "prior.transformer._attn_mods.64.attn.c_attn.b", "prior.transformer._attn_mods.64.attn.c_proj.w", "prior.transformer._attn_mods.64.attn.c_proj.b", "prior.transformer._attn_mods.64.ln_0.weight", "prior.transformer._attn_mods.64.ln_0.bias", "prior.transformer._attn_mods.64.mlp.c_fc.w", "prior.transformer._attn_mods.64.mlp.c_fc.b", "prior.transformer._attn_mods.64.mlp.c_proj.w", "prior.transformer._attn_mods.64.mlp.c_proj.b", "prior.transformer._attn_mods.64.ln_1.weight", "prior.transformer._attn_mods.64.ln_1.bias", "prior.transformer._attn_mods.65.attn.c_attn.w", "prior.transformer._attn_mods.65.attn.c_attn.b", "prior.transformer._attn_mods.65.attn.c_proj.w", "prior.transformer._attn_mods.65.attn.c_proj.b", "prior.transformer._attn_mods.65.ln_0.weight", "prior.transformer._attn_mods.65.ln_0.bias", "prior.transformer._attn_mods.65.mlp.c_fc.w", "prior.transformer._attn_mods.65.mlp.c_fc.b", "prior.transformer._attn_mods.65.mlp.c_proj.w", "prior.transformer._attn_mods.65.mlp.c_proj.b", "prior.transformer._attn_mods.65.ln_1.weight", "prior.transformer._attn_mods.65.ln_1.bias", "prior.transformer._attn_mods.66.attn.c_attn.w", "prior.transformer._attn_mods.66.attn.c_attn.b", "prior.transformer._attn_mods.66.attn.c_proj.w", "prior.transformer._attn_mods.66.attn.c_proj.b", "prior.transformer._attn_mods.66.ln_0.weight", "prior.transformer._attn_mods.66.ln_0.bias", "prior.transformer._attn_mods.66.mlp.c_fc.w", "prior.transformer._attn_mods.66.mlp.c_fc.b", "prior.transformer._attn_mods.66.mlp.c_proj.w", "prior.transformer._attn_mods.66.mlp.c_proj.b", "prior.transformer._attn_mods.66.ln_1.weight", "prior.transformer._attn_mods.66.ln_1.bias", "prior.transformer._attn_mods.67.attn.c_attn.w", "prior.transformer._attn_mods.67.attn.c_attn.b", "prior.transformer._attn_mods.67.attn.c_proj.w", "prior.transformer._attn_mods.67.attn.c_proj.b", "prior.transformer._attn_mods.67.ln_0.weight", "prior.transformer._attn_mods.67.ln_0.bias", "prior.transformer._attn_mods.67.mlp.c_fc.w", "prior.transformer._attn_mods.67.mlp.c_fc.b", "prior.transformer._attn_mods.67.mlp.c_proj.w", "prior.transformer._attn_mods.67.mlp.c_proj.b", "prior.transformer._attn_mods.67.ln_1.weight", "prior.transformer._attn_mods.67.ln_1.bias", "prior.transformer._attn_mods.68.attn.c_attn.w", "prior.transformer._attn_mods.68.attn.c_attn.b", "prior.transformer._attn_mods.68.attn.c_proj.w", "prior.transformer._attn_mods.68.attn.c_proj.b", "prior.transformer._attn_mods.68.ln_0.weight", "prior.transformer._attn_mods.68.ln_0.bias", "prior.transformer._attn_mods.68.mlp.c_fc.w", "prior.transformer._attn_mods.68.mlp.c_fc.b", "prior.transformer._attn_mods.68.mlp.c_proj.w", "prior.transformer._attn_mods.68.mlp.c_proj.b", "prior.transformer._attn_mods.68.ln_1.weight", "prior.transformer._attn_mods.68.ln_1.bias", "prior.transformer._attn_mods.69.attn.c_attn.w", "prior.transformer._attn_mods.69.attn.c_attn.b", "prior.transformer._attn_mods.69.attn.c_proj.w", "prior.transformer._attn_mods.69.attn.c_proj.b", "prior.transformer._attn_mods.69.ln_0.weight", "prior.transformer._attn_mods.69.ln_0.bias", "prior.transformer._attn_mods.69.mlp.c_fc.w", "prior.transformer._attn_mods.69.mlp.c_fc.b", "prior.transformer._attn_mods.69.mlp.c_proj.w", "prior.transformer._attn_mods.69.mlp.c_proj.b", "prior.transformer._attn_mods.69.ln_1.weight", "prior.transformer._attn_mods.69.ln_1.bias", "prior.transformer._attn_mods.70.attn.c_attn.w", "prior.transformer._attn_mods.70.attn.c_attn.b", "prior.transformer._attn_mods.70.attn.c_proj.w", "prior.transformer._attn_mods.70.attn.c_proj.b", "prior.transformer._attn_mods.70.ln_0.weight", "prior.transformer._attn_mods.70.ln_0.bias", "prior.transformer._attn_mods.70.mlp.c_fc.w", "prior.transformer._attn_mods.70.mlp.c_fc.b", "prior.transformer._attn_mods.70.mlp.c_proj.w", "prior.transformer._attn_mods.70.mlp.c_proj.b", "prior.transformer._attn_mods.70.ln_1.weight", "prior.transformer._attn_mods.70.ln_1.bias", "prior.transformer._attn_mods.71.attn.c_attn.w", "prior.transformer._attn_mods.71.attn.c_attn.b", "prior.transformer._attn_mods.71.attn.c_proj.w", "prior.transformer._attn_mods.71.attn.c_proj.b", "prior.transformer._attn_mods.71.ln_0.weight", "prior.transformer._attn_mods.71.ln_0.bias", "prior.transformer._attn_mods.71.mlp.c_fc.w", "prior.transformer._attn_mods.71.mlp.c_fc.b", "prior.transformer._attn_mods.71.mlp.c_proj.w", "prior.transformer._attn_mods.71.mlp.c_proj.b", "prior.transformer._attn_mods.71.ln_1.weight", "prior.transformer._attn_mods.71.ln_1.bias". 

[rodrigo-castellon]

Hi, Thanks for building off of our work! Notice that this patch fixes this issue by ensuring that strict=False when loading in a model (see this PyTorch forum thread for more info about what this does). You need to apply this patch to jukebox/jukebox/make_models.py as done on this line of the Dockerfile.

Hope this helps.

[marypilataki]

Great, thank you very much for your response, this has resolved my issue.

Best wishes