choderalab / mtenn

Modular Training and Evaluation of Neural Networks
MIT License
5 stars 1 forks source link

Connect the individual prediction gradients into the Combination workflow #66

Open kaminow opened 3 weeks ago

kaminow commented 3 weeks ago

Currently the Combination implementations only return the final prediction, so if we want to calculate a loss against the individual pose predictions, we get an error since the computation graph for the individual predictions are empty. Want to connect in the individual pose prediction gradients to allow using a loss against these pose predictions (eg cross entropy loss for selecting the best pose).

kaminow commented 3 weeks ago

this might just be as simple as returning the individual pose predictions in the Combination.forward methods and then returning their gradients in the Combination.backward methods

kaminow commented 6 days ago

@jchodera @hmacdope I think I'm pretty happy with the level of commenting/documentation now, but please take a look when you get a chance and let me know what you think!