checkpoints don't match model

The error I previously had:
RuntimeError: Error(s) in loading state_dict for TrEGNN:
size mismatch for backbone_block.transformer.encoders.0.mha_norm.weight: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([128]).
size mismatch for backbone_block.transformer.encoders.0.mha_norm.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([128]).
size mismatch for backbone_block.transformer.encoders.0.ffn_norm.weight: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([128]).
size mismatch for backbone_block.transformer.encoders.0.ffn_norm.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([128]).
size mismatch for backbone_block.transformer.encoders.1.mha_norm.weight: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([128]).
size mismatch for backbone_block.transformer.encoders.1.mha_norm.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([128]).
size mismatch for backbone_block.transformer.encoders.1.ffn_norm.weight: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([128]).
size mismatch for backbone_block.transformer.encoders.1.ffn_norm.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([128]).
size mismatch for torsion_block.transformer.encoders.0.mha_norm.weight: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([128]).
size mismatch for torsion_block.transformer.encoders.0.mha_norm.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([128]).
size mismatch for torsion_block.transformer.encoders.0.ffn_norm.weight: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([128]).
size mismatch for torsion_block.transformer.encoders.0.ffn_norm.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([128]).
size mismatch for torsion_block.transformer.encoders.1.mha_norm.weight: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([128]).
size mismatch for torsion_block.transformer.encoders.1.mha_norm.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([128]).
size mismatch for torsion_block.transformer.encoders.1.ffn_norm.weight: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([128]).
size mismatch for torsion_block.transformer.encoders.1.ffn_norm.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([128]).

Was resolved by downgrading pytorch geometric.