Missing keys:
go_encoder.gat_layers.0.att_src
go_encoder.gat_layers.0.att_dst
go_encoder.gat_layers.0.bias
go_encoder.gat_layers.0.lin.weight
go_encoder.gat_layers.1.att_src
go_encoder.gat_layers.1.att_dst
go_encoder.gat_layers.1.bias
go_encoder.gat_layers.1.lin.weight
go_encoder.gat_layers.2.att_src
go_encoder.gat_layers.2.att_dst
go_encoder.gat_layers.2.bias
go_encoder.gat_layers.2.lin.weight
go_encoder.output_projection.0.weight
go_encoder.output_projection.0.bias
go_encoder.output_projection.2.weight
go_encoder.output_projection.2.bias
go_encoder.cross_attention_reducer.query_embeddings
go_encoder.cross_attention_reducer.cross_attention.in_proj_weight
go_encoder.cross_attention_reducer.cross_attention.in_proj_bias
go_encoder.cross_attention_reducer.cross_attention.out_proj.weight
go_encoder.cross_attention_reducer.cross_attention.out_proj.bias
go_encoder.cross_attention_reducer.layer_norm1.weight
go_encoder.cross_attention_reducer.layer_norm1.bias
go_encoder.cross_attention_reducer.layer_norm2.weight
go_encoder.cross_attention_reducer.layer_norm2.bias
go_encoder.cross_attention_reducer.ffn.0.weight
go_encoder.cross_attention_reducer.ffn.0.bias
go_encoder.cross_attention_reducer.ffn.3.weight
go_encoder.cross_attention_reducer.ffn.3.bias

Unexpected keys:
protein_model._structure_encoder.transformer.blocks.0.geom_attn.distance_scale_per_head
protein_model._structure_encoder.transformer.blocks.0.geom_attn.rotation_scale_per_head
protein_model._structure_encoder.transformer.blocks.0.geom_attn.s_norm.weight
protein_model._structure_encoder.transformer.blocks.0.geom_attn.s_norm.bias
protein_model._structure_encoder.transformer.blocks.0.geom_attn.proj.weight
protein_model._structure_encoder.transformer.blocks.0.geom_attn.proj.bias
protein_model._structure_encoder.transformer.blocks.0.geom_attn.out_proj.weight
protein_model._structure_encoder.transformer.blocks.0.geom_attn.out_proj.bias
protein_model._structure_encoder.transformer.blocks.0.ffn.0.weight
protein_model._structure_encoder.transformer.blocks.0.ffn.0.bias
protein_model._structure_encoder.transformer.blocks.0.ffn.1.weight
protein_model._structure_encoder.transformer.blocks.0.ffn.1.bias
protein_model._structure_encoder.transformer.blocks.0.ffn.3.weight
protein_model._structure_encoder.transformer.blocks.0.ffn.3.bias
protein_model._structure_encoder.transformer.blocks.1.geom_attn.distance_scale_per_head
protein_model._structure_encoder.transformer.blocks.1.geom_attn.rotation_scale_per_head
protein_model._structure_encoder.transformer.blocks.1.geom_attn.s_norm.weight
protein_model._structure_encoder.transformer.blocks.1.geom_attn.s_norm.bias
protein_model._structure_encoder.transformer.blocks.1.geom_attn.proj.weight
protein_model._structure_encoder.transformer.blocks.1.geom_attn.proj.bias
protein_model._structure_encoder.transformer.blocks.1.geom_attn.out_proj.weight
protein_model._structure_encoder.transformer.blocks.1.geom_attn.out_proj.bias
protein_model._structure_encoder.transformer.blocks.1.ffn.0.weight
protein_model._structure_encoder.transformer.blocks.1.ffn.0.bias
protein_model._structure_encoder.transformer.blocks.1.ffn.1.weight
protein_model._structure_encoder.transformer.blocks.1.ffn.1.bias
protein_model._structure_encoder.transformer.blocks.1.ffn.3.weight
protein_model._structure_encoder.transformer.blocks.1.ffn.3.bias
protein_model._structure_encoder.pre_vq_proj.weight
protein_model._structure_encoder.pre_vq_proj.bias
protein_model._structure_encoder.codebook.embeddings
protein_model._structure_encoder.codebook.N
protein_model._structure_encoder.codebook.z_avg
protein_model._structure_encoder.relative_positional_embedding.embedding.weight