Dear authors:
Thank you for this nice work.
When I try to load the checkpoint you provided for image sampling of COCO dataset, I got this error. It seems that number of classes is 150 rather than 183.
RuntimeError: Error(s) in loading state_dict for UNetModel:
size mismatch for middle_block.0.in_norm.mlp_shared.0.weight: copying a param with shape torch.Size([128, 151, 3, 3]) from checkpoint, the shape in current model is torch.Size([128, 184, 3, 3]).
size mismatch for middle_block.0.out_norm.mlp_shared.0.weight: copying a param with shape torch.Size([128, 151, 3, 3]) from checkpoint, the shape in current model is torch.Size([128, 184, 3, 3]).
size mismatch for middle_block.2.in_norm.mlp_shared.0.weight: copying a param with shape torch.Size([128, 151, 3, 3]) from checkpoint, the shape in current model is torch.Size([128, 184, 3, 3]).
size mismatch for middle_block.2.out_norm.mlp_shared.0.weight: copying a param with shape torch.Size([128, 151, 3, 3]) from checkpoint, the shape in current model is torch.Size([128, 184, 3, 3]).
size mismatch for output_blocks.0.0.in_norm.mlp_shared.0.weight: copying a param with shape torch.Size([128, 151, 3, 3]) from checkpoint, the shape in current model is torch.Size([128, 184, 3, 3]).
size mismatch for output_blocks.0.0.out_norm.mlp_shared.0.weight: copying a param with shape torch.Size([128, 151, 3, 3]) from checkpoint, the shape in current model is torch.Size([128, 184, 3, 3]).
...
Dear authors: Thank you for this nice work. When I try to load the checkpoint you provided for image sampling of COCO dataset, I got this error. It seems that number of classes is 150 rather than 183. RuntimeError: Error(s) in loading state_dict for UNetModel: size mismatch for middle_block.0.in_norm.mlp_shared.0.weight: copying a param with shape torch.Size([128, 151, 3, 3]) from checkpoint, the shape in current model is torch.Size([128, 184, 3, 3]). size mismatch for middle_block.0.out_norm.mlp_shared.0.weight: copying a param with shape torch.Size([128, 151, 3, 3]) from checkpoint, the shape in current model is torch.Size([128, 184, 3, 3]). size mismatch for middle_block.2.in_norm.mlp_shared.0.weight: copying a param with shape torch.Size([128, 151, 3, 3]) from checkpoint, the shape in current model is torch.Size([128, 184, 3, 3]). size mismatch for middle_block.2.out_norm.mlp_shared.0.weight: copying a param with shape torch.Size([128, 151, 3, 3]) from checkpoint, the shape in current model is torch.Size([128, 184, 3, 3]). size mismatch for output_blocks.0.0.in_norm.mlp_shared.0.weight: copying a param with shape torch.Size([128, 151, 3, 3]) from checkpoint, the shape in current model is torch.Size([128, 184, 3, 3]). size mismatch for output_blocks.0.0.out_norm.mlp_shared.0.weight: copying a param with shape torch.Size([128, 151, 3, 3]) from checkpoint, the shape in current model is torch.Size([128, 184, 3, 3]). ...