Simplify computation of strain (#260)

* Update ase.py

* Update ase.py

torchani/ase.py

@@ -55,13 +55,9 @@ class Calculator(ase.calculators.calculator.Calculator):
 
     @staticmethod
     def strain(tensor, displacement, surface_normal_axis):
-        rest_axes = {0, 1, 2} - set([surface_normal_axis])
-        displacement_normal = displacement[surface_normal_axis]
         displacement_of_tensor = torch.zeros_like(tensor)
-        displacement_of_tensor[..., surface_normal_axis] = tensor[..., surface_normal_axis] * displacement_normal
-        for axis in rest_axes:
-            displacement_axis = displacement[axis]
-            displacement_of_tensor[..., axis] = tensor[..., surface_normal_axis] * displacement_axis
+        for axis in range(3):
+            displacement_of_tensor[..., axis] = tensor[..., surface_normal_axis] * displacement[axis]
         return displacement_of_tensor
 
     def calculate(self, atoms=None, properties=['energy'],