Unverified Commit 077c12b8 authored by Gao, Xiang's avatar Gao, Xiang Committed by GitHub
Browse files

Use the new torch.unique_consecutive to speedup (#209)

parent 276bc75b
...@@ -237,7 +237,7 @@ def triple_by_molecule(atom_index1, atom_index2): ...@@ -237,7 +237,7 @@ def triple_by_molecule(atom_index1, atom_index2):
sorted_ai1, rev_indices = ai1.sort() sorted_ai1, rev_indices = ai1.sort()
# sort and compute unique key # sort and compute unique key
uniqued_central_atom_index, _, counts = torch._unique2_temporary_will_remove_soon(sorted_ai1, sorted=True, return_inverse=False, return_counts=True) uniqued_central_atom_index, counts = torch.unique_consecutive(sorted_ai1, return_counts=True)
# do local combinations within unique key, assuming sorted # do local combinations within unique key, assuming sorted
pair_sizes = counts * (counts - 1) // 2 pair_sizes = counts * (counts - 1) // 2
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment