Fix force shape when only 1 atom (#409)

aiqm · Jan 27, 2020 · ca47b22 · ca47b22
1 parent 987711c
commit ca47b22
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchani/ase.py b/torchani/ase.py
@@ -73,7 +73,7 @@ def calculate(self, atoms=None, properties=['energy'],
 
         if 'forces' in properties:
             forces = -torch.autograd.grad(energy.squeeze(), coordinates)[0]
-            self.results['forces'] = forces.squeeze().to('cpu').numpy()
+            self.results['forces'] = forces.squeeze(0).to('cpu').numpy()
 
         if 'stress' in properties:
             volume = self.atoms.get_volume()