Making MACE using cuEquivariance work with Torch Script #871
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
I have an application where I need to use TorchScript, and noticed that it wasn't working with cuEquivariance in the latest version. I have isolated it down to the two instances in the code where a "self.cueq_config" variable is set, with TorchScript not being compatible with the CuEquivarianceConfig class. I have replaced the "self.cueq_config" values with flags taken from cueq_config. A small working example is attached that is runnable with "python test.py --xyz_file=carbon.xyz", which should fail in the current version of MACE but work with the proposed changes.
test_mace_script.tar.gz