Errors when trying to engage with Jukebox code

[Rsalganik1123]

As the title says, I run into errors when I try to run your code for jukebox/main.py.

In particular, the line:

top_prior = make_prior(hparams, vqvae, device)

Triggers the following error:

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",... 

Would really appreciate advice on fix since it makes training your model impossible.

Thanks in advance!