Tiny improvement (#205)

torchani/aev.py

@@ -258,11 +258,9 @@ def triple_by_molecule(atom_index1, atom_index2):
     local_index2 = rev_indices[sorted_local_index2]
 
     # compute mapping between representation of central-other to pair
-    sign1 = torch.where(local_index1 < n, torch.ones_like(local_index1), -torch.ones_like(local_index1))
-    sign2 = torch.where(local_index2 < n, torch.ones_like(local_index2), -torch.ones_like(local_index2))
-    pair_index1 = torch.where(local_index1 < n, local_index1, local_index1 - n)
-    pair_index2 = torch.where(local_index2 < n, local_index2, local_index2 - n)
-    return central_atom_index, pair_index1, pair_index2, sign1, sign2
+    sign1 = ((local_index1 < n) * 2).to(torch.long) - 1
+    sign2 = ((local_index2 < n) * 2).to(torch.long) - 1
+    return central_atom_index, local_index1 % n, local_index2 % n, sign1, sign2
 
 
 # torch.jit.script