[Bug]:MatterSimCalculator.__setstate__ ignores saved model architecture, causing copy() to rebuild the wrong-sized model

Contact Details

ella.bolland@trinity.ox.ac.uk

Bug Description

Copying a MatterSimCalculator loaded from a non-default checkpoint (e.g. the 5M model) via Python's standard copy.copy() crashes with RuntimeError: size mismatch, because __setstate__ discards the saved model architecture (model_args) and always rebuilds using the hardcoded default (1M/128-unit) architecture instead, regardless of which checkpoint was actually loaded. It then tries to load the real (larger) state dict into that wrong-sized model.

MatterSim Version

1.2.4

Python Version

3.12.13

Reproduction Steps

1. Load a MatterSim calculator from a non-default checkpoint (the 5M model; the 1M model is unaffected, since it happens to match the hardcoded default from_checkpoint falls back to):

from mattersim.forcefield import MatterSimCalculator
calc = MatterSimCalculator(load_path="mattersim-v1.0.0-5m", device="cpu")

2. Copy the calculator using Python's standard library:

from copy import copy
copy(calc)

3. This raises RuntimeError: Error(s) in loading state_dict for M3Gnet: size mismatch for edge_encoder.mlp.0.linear.weight: copying a param with shape torch.Size([256, 4]) from checkpoint, the shape in current model is torch.Size([128, 4]). (and many more identically-shaped mismatches, all 256-vs-128).

Cause: setstate pops model_args from the pickled state but never uses it - always rebuilds via Potential.from_checkpoint(device=self.device) with no load_path (defaults to 1M), then loads the real state dict into that wrong-sized model.

Likely introduced by #153

Expected Behavior

copy.copy() on a MatterSimCalculator should produce an independent calculator equivalent to the original: same architecture, same weights, same predictions, regardless of which checkpoint was loaded.

Actual Behavior

Raises RuntimeError: Error(s) in loading state_dict for M3Gnet, with a long list of size mismatches (256 vs 128 units) between the checkpoint's weights and the rebuilt model. The copy is never produced; the calculator is unusable after copying.

Error Logs

Traceback (most recent call last):
  File "<string>", line 5, in <module>
  File ".../copy.py", line 97, in copy
    return _reconstruct(x, None, *rv)
  File ".../copy.py", line 261, in _reconstruct
    y.__setstate__(state)
  File ".../mattersim/forcefield/potential.py", line 1270, in __setstate__
    self.potential.model.load_state_dict(model_state_dict)
  File ".../torch/nn/modules/module.py", line 2624, in load_state_dict
    raise RuntimeError(
RuntimeError: Error(s) in loading state_dict for M3Gnet:
	Unexpected key(s) in state_dict: "graph_conv.3.gated_mlp_atom.g.0.linear.weight", "graph_conv.3.gated_mlp_atom.g.0.linear.bias", "graph_conv.3.gated_mlp_atom.g.1.linear.weight", "graph_conv.3.gated_mlp_atom.g.1.linear.bias", "graph_conv.3.gated_mlp_atom.sigma.0.linear.weight", "graph_conv.3.gated_mlp_atom.sigma.0.linear.bias", "graph_conv.3.gated_mlp_atom.sigma.1.linear.weight", "graph_conv.3.gated_mlp_atom.sigma.1.linear.bias", "graph_conv.3.edge_layer_atom.linear.weight", "graph_conv.3.gated_mlp_edge.g.0.linear.weight", "graph_conv.3.gated_mlp_edge.g.0.linear.bias", "graph_conv.3.gated_mlp_edge.g.1.linear.weight", "graph_conv.3.gated_mlp_edge.g.1.linear.bias", "graph_conv.3.gated_mlp_edge.sigma.0.linear.weight", "graph_conv.3.gated_mlp_edge.sigma.0.linear.bias", "graph_conv.3.gated_mlp_edge.sigma.1.linear.weight", "graph_conv.3.gated_mlp_edge.sigma.1.linear.bias", "graph_conv.3.edge_layer_edge.linear.weight", "graph_conv.3.three_body.atom_mlp.linear.weight", "graph_conv.3.three_body.atom_mlp.linear.bias", "graph_conv.3.three_body.edge_gate_mlp.g.0.linear.weight", "graph_conv.3.three_body.edge_gate_mlp.sigma.0.linear.weight". 
	size mismatch for edge_encoder.mlp.0.linear.weight: copying a param with shape torch.Size([256, 4]) from checkpoint, the shape in current model is torch.Size([128, 4]).
	size mismatch for graph_conv.0.gated_mlp_atom.g.0.linear.weight: copying a param with shape torch.Size([256, 768]) from checkpoint, the shape in current model is torch.Size([128, 384]).
	size mismatch for graph_conv.0.gated_mlp_atom.g.0.linear.bias: copying a param with shape torch.Size([256]) from checkpoint, the shape in current model is torch.Size([128]).
	size mismatch for graph_conv.0.gated_mlp_atom.g.1.linear.weight: copying a param with shape torch.Size([256, 256]) from checkpoint, the shape in current model is torch.Size([128, 128]).
	size mismatch for graph_conv.0.gated_mlp_atom.g.1.linear.bias: copying a param with shape torch.Size([256]) from checkpoint, the shape in current model is torch.Size([128]).
	size mismatch for graph_conv.0.gated_mlp_atom.sigma.0.linear.weight: copying a param with shape torch.Size([256, 768]) from checkpoint, the shape in current model is torch.Size([128, 384]).
	size mismatch for graph_conv.0.gated_mlp_atom.sigma.0.linear.bias: copying a param with shape torch.Size([256]) from checkpoint, the shape in current model is torch.Size([128]).
	size mismatch for graph_conv.0.gated_mlp_atom.sigma.1.linear.weight: copying a param with shape torch.Size([256, 256]) from checkpoint, the shape in current model is torch.Size([128, 128]).
	size mismatch for graph_conv.0.gated_mlp_atom.sigma.1.linear.bias: copying a param with shape torch.Size([256]) from checkpoint, the shape in current model is torch.Size([128]).
	size mismatch for graph_conv.0.edge_layer_atom.linear.weight: copying a param with shape torch.Size([256, 4]) from checkpoint, the shape in current model is torch.Size([128, 4]).
	size mismatch for graph_conv.0.gated_mlp_edge.g.0.linear.weight: copying a param with shape torch.Size([256, 768]) from checkpoint, the shape in current model is torch.Size([128, 384]).
	size mismatch for graph_conv.0.gated_mlp_edge.g.0.linear.bias: copying a param with shape torch.Size([256]) from checkpoint, the shape in current model is torch.Size([128]).
	size mismatch for graph_conv.0.gated_mlp_edge.g.1.linear.weight: copying a param with shape torch.Size([256, 256]) from checkpoint, the shape in current model is torch.Size([128, 128]).
	size mismatch for graph_conv.0.gated_mlp_edge.g.1.linear.bias: copying a param with shape torch.Size([256]) from checkpoint, the shape in current model is torch.Size([128]).
	size mismatch for graph_conv.0.gated_mlp_edge.sigma.0.linear.weight: copying a param with shape torch.Size([256, 768]) from checkpoint, the shape in current model is torch.Size([128, 384]).
	size mismatch for graph_conv.0.gated_mlp_edge.sigma.0.linear.bias: copying a param with shape torch.Size([256]) from checkpoint, the shape in current model is torch.Size([128]).
	size mismatch for graph_conv.0.gated_mlp_edge.sigma.1.linear.weight: copying a param with shape torch.Size([256, 256]) from checkpoint, the shape in current model is torch.Size([128, 128]).
	size mismatch for graph_conv.0.gated_mlp_edge.sigma.1.linear.bias: copying a param with shape torch.Size([256]) from checkpoint, the shape in current model is torch.Size([128]).
	size mismatch for graph_conv.0.edge_layer_edge.linear.weight: copying a param with shape torch.Size([256, 4]) from checkpoint, the shape in current model is torch.Size([128, 4]).
	size mismatch for graph_conv.0.three_body.atom_mlp.linear.weight: copying a param with shape torch.Size([16, 256]) from checkpoint, the shape in current model is torch.Size([16, 128]).
	size mismatch for graph_conv.0.three_body.edge_gate_mlp.g.0.linear.weight: copying a param with shape torch.Size([256, 16]) from checkpoint, the shape in current model is torch.Size([128, 16]).
	size mismatch for graph_conv.0.three_body.edge_gate_mlp.sigma.0.linear.weight: copying a param with shape torch.Size([256, 16]) from checkpoint, the shape in current model is torch.Size([128, 16]).
	size mismatch for final.g.0.linear.weight: copying a param with shape torch.Size([256, 256]) from checkpoint, the shape in current model is torch.Size([128, 128]).
	size mismatch for final.g.0.linear.bias: copying a param with shape torch.Size([256]) from checkpoint, the shape in current model is torch.Size([128]).
	size mismatch for final.g.1.linear.weight: copying a param with shape torch.Size([256, 256]) from checkpoint, the shape in current model is torch.Size([128, 128]).
	size mismatch for final.g.1.linear.bias: copying a param with shape torch.Size([256]) from checkpoint, the shape in current model is torch.Size([128]).
	size mismatch for final.g.2.linear.weight: copying a param with shape torch.Size([1, 256]) from checkpoint, the shape in current model is torch.Size([1, 128]).
	size mismatch for final.sigma.0.linear.weight: copying a param with shape torch.Size([256, 256]) from checkpoint, the shape in current model is torch.Size([128, 128]).
	size mismatch for final.sigma.0.linear.bias: copying a param with shape torch.Size([256]) from checkpoint, the shape in current model is torch.Size([128]).
	size mismatch for final.sigma.1.linear.weight: copying a param with shape torch.Size([256, 256]) from checkpoint, the shape in current model is torch.Size([128, 128]).
	size mismatch for final.sigma.1.linear.bias: copying a param with shape torch.Size([256]) from checkpoint, the shape in current model is torch.Size([128]).
	size mismatch for final.sigma.2.linear.weight: copying a param with shape torch.Size([1, 256]) from checkpoint, the shape in current model is torch.Size([1, 128]).
	size mismatch for atom_embedding.mlp.0.linear.weight: copying a param with shape torch.Size([256, 95]) from checkpoint, the shape in current model is torch.Size([128, 95]).

Code of Conduct

• I agree to follow the project's Code of Conduct

Start at mattersim/forcefield/potential.py around setstate at line 1270, then trace how MatterSimCalculator stores model_args and rebuilds its Potential during copy.copy(). Reproduce the issue with the 5M checkpoint on CPU and verify that copying completes with the original architecture, weights, and predictions preserved.