Size mismatch when loading 60M no tri checkpoint

[WesleyHsieh0806]

Hi,

I tried to load the 60M no tri model and encountered the error messages below.
There seems to be a misalignment between the config and the checkpoint.

RuntimeError: Error(s) in loading state_dict for Proteina:
        size mismatch for nn.init_repr_factory.linear_out.weight: copying a param with shape torch.Size([512, 200]) from checkpoint, the shape in current model is torch.Size([512, 132]).
        size mismatch for nn.cond_factory.feat_creators.1.embedding_C.weight: copying a param with shape torch.Size([6, 196]) from checkpoint, the shape in current model is torch.Size([6, 256]).
        size mismatch for nn.cond_factory.feat_creators.1.embedding_A.weight: copying a param with shape torch.Size([44, 196]) from checkpoint, the shape in current model is torch.Size([44, 256]).
        size mismatch for nn.cond_factory.feat_creators.1.embedding_T.weight: copying a param with shape torch.Size([1473, 196]) from checkpoint, the shape in current model is torch.Size([1473, 256]).
        size mismatch for nn.cond_factory.linear_out.weight: copying a param with shape torch.Size([128, 784]) from checkpoint, the shape in current model is torch.Size([512, 1024]).
        size mismatch for nn.transition_c_1.swish_linear.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([2048, 512]).
        size mismatch for nn.transition_c_1.linear_out.weight: copying a param with shape torch.Size([128, 256]) from checkpoint, the shape in current model is torch.Size([512, 1024]).
        size mismatch for nn.transition_c_2.swish_linear.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([2048, 512]).
        size mismatch for nn.transition_c_2.linear_out.weight: copying a param with shape torch.Size([128, 256]) from checkpoint, the shape in current model is torch.Size([512, 1024]).
        size mismatch for nn.pair_repr_builder.init_repr_factory.ln_out.weight: copying a param with shape torch.Size([196]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for nn.pair_repr_builder.init_repr_factory.ln_out.bias: copying a param with shape torch.Size([196]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for nn.pair_repr_builder.init_repr_factory.linear_out.weight: copying a param with shape torch.Size([196, 319]) from checkpoint, the shape in current model is torch.Size([256, 319]).
        size mismatch for nn.pair_repr_builder.cond_factory.ln_out.weight: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.pair_repr_builder.cond_factory.ln_out.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.pair_repr_builder.cond_factory.linear_out.weight: copying a param with shape torch.Size([128, 196]) from checkpoint, the shape in current model is torch.Size([512, 256]).
        size mismatch for nn.pair_repr_builder.adaln.norm_cond.weight: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.pair_repr_builder.adaln.norm_cond.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.pair_repr_builder.adaln.to_gamma.0.weight: copying a param with shape torch.Size([196, 128]) from checkpoint, the shape in current model is torch.Size([256, 512]).
        size mismatch for nn.pair_repr_builder.adaln.to_gamma.0.bias: copying a param with shape torch.Size([196]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for nn.pair_repr_builder.adaln.to_beta.weight: copying a param with shape torch.Size([196, 128]) from checkpoint, the shape in current model is torch.Size([256, 512]).
        size mismatch for nn.transformer_layers.0.mhba.adaln.norm_cond.weight: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.0.mhba.adaln.norm_cond.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.0.mhba.adaln.to_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.0.mhba.adaln.to_beta.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.0.mhba.mha.to_qkv.weight: copying a param with shape torch.Size([1512, 512]) from checkpoint, the shape in current model is torch.Size([1536, 512]).
        size mismatch for nn.transformer_layers.0.mhba.mha.to_qkv.bias: copying a param with shape torch.Size([1512]) from checkpoint, the shape in current model is torch.Size([1536]).
        size mismatch for nn.transformer_layers.0.mhba.mha.to_g.weight: copying a param with shape torch.Size([504, 512]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.0.mhba.mha.to_g.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.0.mhba.mha.to_out_node.weight: copying a param with shape torch.Size([512, 504]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.0.mhba.mha.q_layer_norm.weight: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.0.mhba.mha.q_layer_norm.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.0.mhba.mha.k_layer_norm.weight: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.0.mhba.mha.k_layer_norm.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.0.mhba.mha.to_bias.weight: copying a param with shape torch.Size([12, 196]) from checkpoint, the shape in current model is torch.Size([8, 256]).
        size mismatch for nn.transformer_layers.0.mhba.mha.pair_norm.weight: copying a param with shape torch.Size([196]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for nn.transformer_layers.0.mhba.mha.pair_norm.bias: copying a param with shape torch.Size([196]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for nn.transformer_layers.0.mhba.scale_output.to_adaln_zero_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.0.transition.adaln.norm_cond.weight: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.0.transition.adaln.norm_cond.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.0.transition.adaln.to_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.0.transition.adaln.to_beta.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.0.transition.scale_output.to_adaln_zero_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.1.mhba.adaln.norm_cond.weight: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.1.mhba.adaln.norm_cond.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.1.mhba.adaln.to_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.1.mhba.adaln.to_beta.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.1.mhba.mha.to_qkv.weight: copying a param with shape torch.Size([1512, 512]) from checkpoint, the shape in current model is torch.Size([1536, 512]).
        size mismatch for nn.transformer_layers.1.mhba.mha.to_qkv.bias: copying a param with shape torch.Size([1512]) from checkpoint, the shape in current model is torch.Size([1536]).
        size mismatch for nn.transformer_layers.1.mhba.mha.to_g.weight: copying a param with shape torch.Size([504, 512]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.1.mhba.mha.to_g.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.1.mhba.mha.to_out_node.weight: copying a param with shape torch.Size([512, 504]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.1.mhba.mha.q_layer_norm.weight: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.1.mhba.mha.q_layer_norm.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.1.mhba.mha.k_layer_norm.weight: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.1.mhba.mha.k_layer_norm.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.1.mhba.mha.to_bias.weight: copying a param with shape torch.Size([12, 196]) from checkpoint, the shape in current model is torch.Size([8, 256]).
        size mismatch for nn.transformer_layers.1.mhba.mha.pair_norm.weight: copying a param with shape torch.Size([196]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for nn.transformer_layers.1.mhba.mha.pair_norm.bias: copying a param with shape torch.Size([196]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for nn.transformer_layers.1.mhba.scale_output.to_adaln_zero_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.1.transition.adaln.norm_cond.weight: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.1.transition.adaln.norm_cond.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.1.transition.adaln.to_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.1.transition.adaln.to_beta.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.1.transition.scale_output.to_adaln_zero_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.2.mhba.adaln.norm_cond.weight: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.2.mhba.adaln.norm_cond.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.2.mhba.adaln.to_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.2.mhba.adaln.to_beta.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.2.mhba.mha.to_qkv.weight: copying a param with shape torch.Size([1512, 512]) from checkpoint, the shape in current model is torch.Size([1536, 512]).
        size mismatch for nn.transformer_layers.2.mhba.mha.to_qkv.bias: copying a param with shape torch.Size([1512]) from checkpoint, the shape in current model is torch.Size([1536]).
        size mismatch for nn.transformer_layers.2.mhba.mha.to_g.weight: copying a param with shape torch.Size([504, 512]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.2.mhba.mha.to_g.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.2.mhba.mha.to_out_node.weight: copying a param with shape torch.Size([512, 504]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.2.mhba.mha.q_layer_norm.weight: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.2.mhba.mha.q_layer_norm.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.2.mhba.mha.k_layer_norm.weight: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.2.mhba.mha.k_layer_norm.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.2.mhba.mha.to_bias.weight: copying a param with shape torch.Size([12, 196]) from checkpoint, the shape in current model is torch.Size([8, 256]).
        size mismatch for nn.transformer_layers.2.mhba.mha.pair_norm.weight: copying a param with shape torch.Size([196]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for nn.transformer_layers.2.mhba.mha.pair_norm.bias: copying a param with shape torch.Size([196]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for nn.transformer_layers.2.mhba.scale_output.to_adaln_zero_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.2.transition.adaln.norm_cond.weight: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.2.transition.adaln.norm_cond.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.2.transition.adaln.to_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.2.transition.adaln.to_beta.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.2.transition.scale_output.to_adaln_zero_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.3.mhba.adaln.norm_cond.weight: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.3.mhba.adaln.norm_cond.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.3.mhba.adaln.to_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.3.mhba.adaln.to_beta.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.3.mhba.mha.to_qkv.weight: copying a param with shape torch.Size([1512, 512]) from checkpoint, the shape in current model is torch.Size([1536, 512]).
        size mismatch for nn.transformer_layers.3.mhba.mha.to_qkv.bias: copying a param with shape torch.Size([1512]) from checkpoint, the shape in current model is torch.Size([1536]).
        size mismatch for nn.transformer_layers.3.mhba.mha.to_g.weight: copying a param with shape torch.Size([504, 512]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.3.mhba.mha.to_g.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.3.mhba.mha.to_out_node.weight: copying a param with shape torch.Size([512, 504]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.3.mhba.mha.q_layer_norm.weight: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.3.mhba.mha.q_layer_norm.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.3.mhba.mha.k_layer_norm.weight: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.3.mhba.mha.k_layer_norm.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.3.mhba.mha.to_bias.weight: copying a param with shape torch.Size([12, 196]) from checkpoint, the shape in current model is torch.Size([8, 256]).
        size mismatch for nn.transformer_layers.3.mhba.mha.pair_norm.weight: copying a param with shape torch.Size([196]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for nn.transformer_layers.3.mhba.mha.pair_norm.bias: copying a param with shape torch.Size([196]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for nn.transformer_layers.3.mhba.scale_output.to_adaln_zero_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.3.transition.adaln.norm_cond.weight: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.3.transition.adaln.norm_cond.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.3.transition.adaln.to_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.3.transition.adaln.to_beta.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.3.transition.scale_output.to_adaln_zero_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.4.mhba.adaln.norm_cond.weight: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.4.mhba.adaln.norm_cond.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.4.mhba.adaln.to_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.4.mhba.adaln.to_beta.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.4.mhba.mha.to_qkv.weight: copying a param with shape torch.Size([1512, 512]) from checkpoint, the shape in current model is torch.Size([1536, 512]).
        size mismatch for nn.transformer_layers.4.mhba.mha.to_qkv.bias: copying a param with shape torch.Size([1512]) from checkpoint, the shape in current model is torch.Size([1536]).
        size mismatch for nn.transformer_layers.4.mhba.mha.to_g.weight: copying a param with shape torch.Size([504, 512]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.4.mhba.mha.to_g.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.4.mhba.mha.to_out_node.weight: copying a param with shape torch.Size([512, 504]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.4.mhba.mha.q_layer_norm.weight: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.4.mhba.mha.q_layer_norm.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.4.mhba.mha.k_layer_norm.weight: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.4.mhba.mha.k_layer_norm.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.4.mhba.mha.to_bias.weight: copying a param with shape torch.Size([12, 196]) from checkpoint, the shape in current model is torch.Size([8, 256]).
        size mismatch for nn.transformer_layers.4.mhba.mha.pair_norm.weight: copying a param with shape torch.Size([196]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for nn.transformer_layers.4.mhba.mha.pair_norm.bias: copying a param with shape torch.Size([196]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for nn.transformer_layers.4.mhba.scale_output.to_adaln_zero_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.4.transition.adaln.norm_cond.weight: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.4.transition.adaln.norm_cond.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.4.transition.adaln.to_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.4.transition.adaln.to_beta.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.4.transition.scale_output.to_adaln_zero_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.5.mhba.adaln.norm_cond.weight: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.5.mhba.adaln.norm_cond.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.5.mhba.adaln.to_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.5.mhba.adaln.to_beta.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.5.mhba.mha.to_qkv.weight: copying a param with shape torch.Size([1512, 512]) from checkpoint, the shape in current model is torch.Size([1536, 512]).
        size mismatch for nn.transformer_layers.5.mhba.mha.to_qkv.bias: copying a param with shape torch.Size([1512]) from checkpoint, the shape in current model is torch.Size([1536]).
        size mismatch for nn.transformer_layers.5.mhba.mha.to_g.weight: copying a param with shape torch.Size([504, 512]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.5.mhba.mha.to_g.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.5.mhba.mha.to_out_node.weight: copying a param with shape torch.Size([512, 504]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.5.mhba.mha.q_layer_norm.weight: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.5.mhba.mha.q_layer_norm.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.5.mhba.mha.k_layer_norm.weight: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.5.mhba.mha.k_layer_norm.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.5.mhba.mha.to_bias.weight: copying a param with shape torch.Size([12, 196]) from checkpoint, the shape in current model is torch.Size([8, 256]).
        size mismatch for nn.transformer_layers.5.mhba.mha.pair_norm.weight: copying a param with shape torch.Size([196]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for nn.transformer_layers.5.mhba.mha.pair_norm.bias: copying a param with shape torch.Size([196]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for nn.transformer_layers.5.mhba.scale_output.to_adaln_zero_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.5.transition.adaln.norm_cond.weight: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.5.transition.adaln.norm_cond.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.5.transition.adaln.to_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.5.transition.adaln.to_beta.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.5.transition.scale_output.to_adaln_zero_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.6.mhba.adaln.norm_cond.weight: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.6.mhba.adaln.norm_cond.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.6.mhba.adaln.to_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.6.mhba.adaln.to_beta.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.6.mhba.mha.to_qkv.weight: copying a param with shape torch.Size([1512, 512]) from checkpoint, the shape in current model is torch.Size([1536, 512]).
        size mismatch for nn.transformer_layers.6.mhba.mha.to_qkv.bias: copying a param with shape torch.Size([1512]) from checkpoint, the shape in current model is torch.Size([1536]).
        size mismatch for nn.transformer_layers.6.mhba.mha.to_g.weight: copying a param with shape torch.Size([504, 512]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.6.mhba.mha.to_g.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.6.mhba.mha.to_out_node.weight: copying a param with shape torch.Size([512, 504]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.6.mhba.mha.q_layer_norm.weight: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.6.mhba.mha.q_layer_norm.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.6.mhba.mha.k_layer_norm.weight: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.6.mhba.mha.k_layer_norm.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.6.mhba.mha.to_bias.weight: copying a param with shape torch.Size([12, 196]) from checkpoint, the shape in current model is torch.Size([8, 256]).
        size mismatch for nn.transformer_layers.6.mhba.mha.pair_norm.weight: copying a param with shape torch.Size([196]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for nn.transformer_layers.6.mhba.mha.pair_norm.bias: copying a param with shape torch.Size([196]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for nn.transformer_layers.6.mhba.scale_output.to_adaln_zero_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.6.transition.adaln.norm_cond.weight: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.6.transition.adaln.norm_cond.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.6.transition.adaln.to_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.6.transition.adaln.to_beta.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.6.transition.scale_output.to_adaln_zero_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.7.mhba.adaln.norm_cond.weight: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.7.mhba.adaln.norm_cond.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.7.mhba.adaln.to_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.7.mhba.adaln.to_beta.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.7.mhba.mha.to_qkv.weight: copying a param with shape torch.Size([1512, 512]) from checkpoint, the shape in current model is torch.Size([1536, 512]).
        size mismatch for nn.transformer_layers.7.mhba.mha.to_qkv.bias: copying a param with shape torch.Size([1512]) from checkpoint, the shape in current model is torch.Size([1536]).
        size mismatch for nn.transformer_layers.7.mhba.mha.to_g.weight: copying a param with shape torch.Size([504, 512]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.7.mhba.mha.to_g.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.7.mhba.mha.to_out_node.weight: copying a param with shape torch.Size([512, 504]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.7.mhba.mha.q_layer_norm.weight: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.7.mhba.mha.q_layer_norm.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.7.mhba.mha.k_layer_norm.weight: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.7.mhba.mha.k_layer_norm.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.7.mhba.mha.to_bias.weight: copying a param with shape torch.Size([12, 196]) from checkpoint, the shape in current model is torch.Size([8, 256]).
        size mismatch for nn.transformer_layers.7.mhba.mha.pair_norm.weight: copying a param with shape torch.Size([196]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for nn.transformer_layers.7.mhba.mha.pair_norm.bias: copying a param with shape torch.Size([196]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for nn.transformer_layers.7.mhba.scale_output.to_adaln_zero_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.7.transition.adaln.norm_cond.weight: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.7.transition.adaln.norm_cond.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.7.transition.adaln.to_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.7.transition.adaln.to_beta.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.7.transition.scale_output.to_adaln_zero_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.8.mhba.adaln.norm_cond.weight: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.8.mhba.adaln.norm_cond.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.8.mhba.adaln.to_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.8.mhba.adaln.to_beta.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.8.mhba.mha.to_qkv.weight: copying a param with shape torch.Size([1512, 512]) from checkpoint, the shape in current model is torch.Size([1536, 512]).
        size mismatch for nn.transformer_layers.8.mhba.mha.to_qkv.bias: copying a param with shape torch.Size([1512]) from checkpoint, the shape in current model is torch.Size([1536]).
        size mismatch for nn.transformer_layers.8.mhba.mha.to_g.weight: copying a param with shape torch.Size([504, 512]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.8.mhba.mha.to_g.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.8.mhba.mha.to_out_node.weight: copying a param with shape torch.Size([512, 504]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.8.mhba.mha.q_layer_norm.weight: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.8.mhba.mha.q_layer_norm.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.8.mhba.mha.k_layer_norm.weight: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.8.mhba.mha.k_layer_norm.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.8.mhba.mha.to_bias.weight: copying a param with shape torch.Size([12, 196]) from checkpoint, the shape in current model is torch.Size([8, 256]).
        size mismatch for nn.transformer_layers.8.mhba.mha.pair_norm.weight: copying a param with shape torch.Size([196]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for nn.transformer_layers.8.mhba.mha.pair_norm.bias: copying a param with shape torch.Size([196]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for nn.transformer_layers.8.mhba.scale_output.to_adaln_zero_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.8.transition.adaln.norm_cond.weight: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.8.transition.adaln.norm_cond.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.8.transition.adaln.to_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.8.transition.adaln.to_beta.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.8.transition.scale_output.to_adaln_zero_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.9.mhba.adaln.norm_cond.weight: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.9.mhba.adaln.norm_cond.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.9.mhba.adaln.to_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.9.mhba.adaln.to_beta.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.9.mhba.mha.to_qkv.weight: copying a param with shape torch.Size([1512, 512]) from checkpoint, the shape in current model is torch.Size([1536, 512]).
        size mismatch for nn.transformer_layers.9.mhba.mha.to_qkv.bias: copying a param with shape torch.Size([1512]) from checkpoint, the shape in current model is torch.Size([1536]).
        size mismatch for nn.transformer_layers.9.mhba.mha.to_g.weight: copying a param with shape torch.Size([504, 512]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.9.mhba.mha.to_g.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.9.mhba.mha.to_out_node.weight: copying a param with shape torch.Size([512, 504]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.9.mhba.mha.q_layer_norm.weight: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.9.mhba.mha.q_layer_norm.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.9.mhba.mha.k_layer_norm.weight: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.9.mhba.mha.k_layer_norm.bias: copying a param with shape torch.Size([504]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.9.mhba.mha.to_bias.weight: copying a param with shape torch.Size([12, 196]) from checkpoint, the shape in current model is torch.Size([8, 256]).
        size mismatch for nn.transformer_layers.9.mhba.mha.pair_norm.weight: copying a param with shape torch.Size([196]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for nn.transformer_layers.9.mhba.mha.pair_norm.bias: copying a param with shape torch.Size([196]) from checkpoint, the shape in current model is torch.Size([256]).
        size mismatch for nn.transformer_layers.9.mhba.scale_output.to_adaln_zero_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.9.transition.adaln.norm_cond.weight: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.9.transition.adaln.norm_cond.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([512]).
        size mismatch for nn.transformer_layers.9.transition.adaln.to_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.9.transition.adaln.to_beta.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).
        size mismatch for nn.transformer_layers.9.transition.scale_output.to_adaln_zero_gamma.0.weight: copying a param with shape torch.Size([512, 128]) from checkpoint, the shape in current model is torch.Size([512, 512]).

[tgeffnerNVIDIA]

Thanks for pointing this out. Can you please try this config?

name: ca_af3  # name of architecture


# Architecture parameters
token_dim: 512  # dimension of the tokens in the sequence
nlayers: 12  # number of transformer layers
nheads: 12  # number of attn heads
residual_mha: True  # whether to use a residual connection in the mha
residual_transition: True  # whether to use a residual connection in the transition
parallel_mha_transition: False  # whether to compute mha and transition as parallel and add them up (AF3 style) or sequentially (normal transofrmers)
use_attn_pair_bias: True  # whether to bias attention using a bias coming from a pair representation

strict_feats: False  # if False, then fills missing features with default values (e.g. chain break with zero, residue sequence index by [0, 1, 2, ...], etc)
# If True, if some feature is not provided, then it raises an error

# List of all available sequence features:
#   "res_seq_pdb_idx", requires transform ResidueSequencePositionPdbTransform
#   "time_emb"
#   "chain_break_per_res", requires transform ChainBreakPerResidueTransform
#   "x_sc"
#   "fold_emb"
feats_init_seq: ["res_seq_pdb_idx", "x_sc", "chain_break_per_res"]  # Sequence features to include in initial representation
feats_cond_seq: ["time_emb", "fold_emb"]  # Sequence features to include in conditioning variables


# Parameters for the features we extract (both for sequence representatoin and conditioning vector)
t_emb_dim: 196  # dimension of the time embedding
dim_cond: 128  # dimension of conditioning vector
idx_emb_dim: 196  # dimension of the sequence position [0, 1, 2, ...] (if contiguous residues) embeddings
fold_emb_dim: 196  # dimension of fold class embedding. This will be multiplied by three, as we have C, A, T embeddings.
cath_code_dir: ${oc.env:DATA_PATH}/pdb_raw/    # This should be set as the path to your pdb_cath dataset directory
multilabel_mode: "sample"


# List of all available sequence features:
#   "xt_pair_dists"
#   "x_sc_pair_dists"
#   "rel_seq_sep"
#   "time_emb"
feats_pair_repr: ["rel_seq_sep", "x_sc_pair_dists", "xt_pair_dists"]  # Features to include in the pair representation
feats_pair_cond: ["time_emb"]  # Features to include in the pair representation conditioning


# Parameters for the pair features we extract
# Binning for the pair distances of noisy xt
xt_pair_dist_dim: 64
xt_pair_dist_min: 0.1  # in nm (not Å)
xt_pair_dist_max: 3  # in nm (not Å)
# Binning for the pair distances for self conditioning
x_sc_pair_dist_dim: 128
x_sc_pair_dist_min: 0.1  # in nm (not Å)
x_sc_pair_dist_max: 3  # in nm (not Å)
# Relative sequence separation
seq_sep_dim: 127  # should be odd >= 5
# Dimension of final pair representation
pair_repr_dim: 196


# Newer stuff for exploration
update_pair_repr: False  # whether to update pair representation, automatically overridden to False if `use_attn_pair_bias: False`
update_pair_repr_every_n: 3  # Update the pair representation every n layers -> For 15 layers we get 5 pair updates (if update pair represnetation is true)
use_tri_mult: False  # whether to use triangular multiplication layers in pair update, ignored if not updating pair representation

num_registers: 10
use_qkln: True

num_buckets_predict_pair: 64