This PR closes #648 by passing cell and pbc as args rather than kwargs when calling self.aev_computer within BuiltinModel and BuiltinEnsemble. This allows a user to add a forward hook to the internal AEVComputer of an ANI model to allow them to get AEVs from the hook at the same time that energies are obtained from the model.
While kwargs work fine with PyTorch, it appears that TorchScript assumes that the input parameter defined in the hook represents all input to the hooked function (AEVComputer.forward) , i.e. both args and kwargs. When recursing to create the scripted version of a module, TorchScript requires that the hooked function is called with args only throughout.
(I have been made aware of the code freeze and appreciate that this won't be merged. Just posting for posterity.)
This PR closes #648 by passing
cellandpbcas args rather than kwargs when callingself.aev_computerwithinBuiltinModelandBuiltinEnsemble. This allows a user to add a forward hook to the internal AEVComputer of an ANI model to allow them to get AEVs from the hook at the same time that energies are obtained from the model.While kwargs work fine with PyTorch, it appears that TorchScript assumes that the
inputparameter defined in the hook represents all input to the hooked function (AEVComputer.forward) , i.e. both args and kwargs. When recursing to create the scripted version of a module, TorchScript requires that the hooked function is called with args only throughout.(I have been made aware of the code freeze and appreciate that this won't be merged. Just posting for posterity.)