Fix neighbor backward #179
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
works Add test for double backwards
Write in full pytorch so autograd works in double gradient Add test using gradcheck and gradgradcheck
|
This is ready for review |
raimis
reviewed
Jun 8, 2023
raimis
reviewed
Jun 8, 2023
raimis
reviewed
Jun 8, 2023
raimis
approved these changes
Jun 14, 2023
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Trying to integrate OptimizedDistance into the equivariant transformer (ET) I realized my module did not supported double backwards (DB), which is required in ET when training with forces.
This PR adds support for double gradient to OptimizedDistance, and adds some tests for it using torch.autograd.gradgradcheck.
After some research I learned about how to support DB in a custom extension. I rewrote the backwards pass in full pytorch to take advantage of autograd for the second gradient. Then I discovered Python custom torch.Autograd.Function are not compatible with jit.script , but it turns out it is compatible if you do it in C++. So I translated the full pytorch backwards to C++. Quite the journey -.-