[davrens101]

Problem

Trainer.save_checkpoint saves the correct best-energy checkpoint to disk, but sets self.best_model = self.model. Later optimizer updates mutate that same object. Consequently, get_best_model() returns the latest weights rather than the weights with the lowest validation energy MAE.

A real CHGNet energy-training step reproduces the discrepancy: the saved best model predicts -3.002830982 eV/atom, while the API returns the later model predicting -2.573733807 eV/atom, after a worse validation epoch.

Change

Snapshot the model only when the best-energy checkpoint is selected. Retrieving the best model does not copy it again or reload a file. Keep the existing validation-selection rule, checkpoint format and public API.

The regression covers subsequent weight updates, a worse validation epoch and a later improvement replacing the stored snapshot.

Verification

CPU, Python 3.12, PyTorch 2.15.0+cpu:

• tests/test_trainer.py: 6 passed, including the new multi-epoch regression and existing training/checkpoint tests (W&B offline).
• Actual CHGNet graph inference, optimizer update and checkpoint smoke: the API best energy now matches the saved best model at -3.002830982 eV/atom, while the latest training model predicts -2.573733807 eV/atom.
• Regression checks both best-state retention after worse epochs and replacement after a later improvement.

No GPU validation is claimed.

Review and assistance

Human review: davrens101 confirmed reviewing the work live before submission. AI assistance was used for investigation, implementation and verification. No maintainer review or approval is claimed.

[codacy-production]

Up to standards ✅

🟢 Issues 0 issues

Results:
0 new issues

View in Codacy

🟢 Metrics 5 complexity

Metric	Results
Complexity	5

View in Codacy

_{NEW Get contextual insights on your PRs based on Codacy's metrics, along with PR and Jira context, without leaving GitHub. Enable AI reviewer
TIP This summary will be updated as you push new changes.}