Remove unnecessary argsort in CUDAHashMap.keys() (#400)

pyg-team · Feb 6, 2025 · 5ec2ff3 · 5ec2ff3
1 parent 8bc4e7a
commit 5ec2ff3
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/pyg_lib/csrc/classes/cuda/hash_map.cu b/pyg_lib/csrc/classes/cuda/hash_map.cu
@@ -84,7 +84,10 @@ struct CUDAHashMapImpl : HashMapImpl {
 
     map_->retrieve_all(key_data, value_data);
 
-    return key.index_select(0, value.argsort());
+    const auto perm = at::empty_like(value);
+    perm[value] = at::arange(value.numel(), value.options());
+
+    return key.index_select(0, perm);
   }
 
  private: