IntelLabs / matsciml

Open MatSci ML Toolkit is a framework for prototyping and scaling out deep learning models for materials discovery supporting widely used materials science datasets, and built on top of PyTorch Lightning, the Deep Graph Library, and PyTorch Geometric.
MIT License
144 stars 20 forks source link

Allow `hidden_irreps` to be passed through the mace wrapper #244

Closed melo-gonzo closed 3 months ago

melo-gonzo commented 3 months ago

There are usecases in which the hidden_irreps may need to be passed in as an o3.Irreps module, such as Irreps("16x0e+16x1o+16x2e"). This update allows the hidden_irreps kwarg to be used in the mace wrapper, and defaults to construct this from atom_embedding_dim if not supplied.