Simplify code using torch.triu_indices (#367)

* Simplify code using torch.triu_indices * Update aev.py
aiqm · Nov 7, 2019 · ae2497d · ae2497d
1 parent 2ba9ae0
commit ae2497d
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/torchani/aev.py b/torchani/aev.py
@@ -168,8 +168,7 @@ def neighbor_pairs(padding_mask, coordinates, cell, shifts, cutoff):
 
 def triu_index(num_species):
     # type: (int) -> torch.Tensor
-    species = torch.arange(num_species, dtype=torch.long)
-    species1, species2 = torch.combinations(species, r=2, with_replacement=True).unbind(-1)
+    species1, species2 = torch.triu_indices(num_species, num_species).unbind(0)
     pair_index = torch.arange(species1.shape[0], dtype=torch.long)
     ret = torch.zeros(num_species, num_species, dtype=torch.long)
     ret[species1, species2] = pair_index