Add retain_graph (force and stress) case in ASE (#417)

This PR resolves #415 Thanks @mhellstr for pointing it out

Add retain_graph (force and stress) case in ASE (#417)
This PR resolves #415 Thanks @mhellstr for pointing it out
00dd8cba · Farhad Ramezanghorbani · GitHub · 79470486 · 00dd8cba
Unverified Commit 00dd8cba authored Feb 05, 2020 by Farhad Ramezanghorbani Committed by GitHub Feb 05, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

torchani/ase.py torchani/ase.py +1 -1

No files found.
--- a/torchani/ase.py
+++ b/torchani/ase.py
@@ -72,7 +72,7 @@ class Calculator(ase.calculators.calculator.Calculator):
        self.results['free_energy'] = energy.item()

        if 'forces' in properties:
-            forces = -torch.autograd.grad(energy.squeeze(), coordinates)[0]
+            forces = -torch.autograd.grad(energy.squeeze(), coordinates, retain_graph='stress' in properties)[0]
            self.results['forces'] = forces.squeeze(0).to('cpu').numpy()

        if 'stress' in properties: