update mace wrapper to enable training

[ys-teh]

ALCHEMI Toolkit Pull Request

Description

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Performance improvement
• Documentation update
• Refactoring (no functional changes)
• CI/CD or infrastructure change

Related Issues

Changes Made

• Changed training=False to training=self.training in MACEWrapper using the PyTorch's flag to be compatible with the upcoming model training workflow.
• Changed mock model used for MACE in unit test to be dependent on a trainable parameter so train mode test makes sense.
• Added a unit test to make sure that optimization now works when using MACEWrapper.

Testing

• Unit tests pass locally (make pytest)
• Linting passes (make lint)
• New tests added for new functionality meets coverage expectations?

Checklist

• I have read and understand the Contributing Guidelines
• I have updated the CHANGELOG.md
• I have performed a self-review of my code
• I have added docstrings to new functions/classes
• I have updated the documentation (if applicable)

Additional Notes

Tip

This repository uses Greptile, an AI code review service, to help conduct
pull request reviews. We encourage contributors to read and consider suggestions
made by Greptile, but note that human maintainers will provide the necessary
reviews for merging: Greptile's comments are not a qualitative judgement
of your code, nor is it an indication that the PR will be accepted/rejected.
We encourage the use of emoji reactions to Greptile comments, depending on
their usefulness and accuracy.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR enables training support in MACEWrapper by replacing the hardcoded training=False with self.training in the forward method, and updates MockMACEModel to wire a learnable parameter through its forward pass so the new optimizer round-trip test is meaningful.

• forward now forwards self.training to the inner MACE model, making dropout/BN-style layers respect the wrapper's current mode during fine-tuning.
• compute_embeddings still passes training=False unconditionally — this is a latent inconsistency if the wrapper is used in train mode for embedding-based training pipelines.
• The new test_train_mode_works_with_optimizer test verifies end-to-end that a gradient step updates the model parameter when the wrapper is in train mode.

Important Files Changed

Filename	Overview
nvalchemi/models/mace.py	Propagates self.training to the inner MACE model in forward; compute_embeddings still hardcodes training=False, which is inconsistent. The from_checkpoint(compile_model=True) path returns a wrapper in train mode despite calling model.eval() internally (flagged in a prior review).
test/models/test_mace.py	Adds test_train_mode_works_with_optimizer to verify gradients flow through the wrapper; updates MockMACEModel to wire a learnable parameter into the forward computation so the test is meaningful.

_{Reviews (2): Last reviewed commit: "add mace train mode test" | Re-trigger Greptile}

[dallasfoster]

/ok to test 1397531

[laserkelvin]

LGTM