Error loading ckpt_mnist.pth: size mismatch for class_embedding.weight

[sweetice100]

I downloaded the pretrained weights (ckpt_mnist.pth) from the Google Drive link in README.

When I try to run:
python test_example_mnist.py --task snapshot_optical --num_classes 10

I get this error:

size mismatch for class_embedding.weight: copying a param with shape torch.Size([10, 512]) from checkpoint, the shape in current model is torch.Size([10, 128]).
size mismatch for l1.0.weight: copying a param with shape torch.Size([1024, 1536]) from checkpoint, the shape in current model is torch.Size([1024, 1152]).

I also tried:

• --task multicolor_optical
• --task iterative_optical
• test_example_celeba.py with all three tasks

All give the same dimension mismatch error.

Could you please tell me the correct command to use these pretrained weights? Or are the weights for an older version of the code?

Thanks!

[TanGeeGo]

Hi, thanks for finding this problem. Can you change your code and make the "dim_expand_ratio=512"?
In this way the code should run and give the results.

I will update code later to fix this issue.

[lyrehcic]

I encountered the same problem, by changing all the "dim_expand_ratio" to dim_expand_ratio=512 can fix the problem and get the result