[Bug] Fix PinSAGE crashes on hypersparse graphs (#1565)

d2e1cfce · Quan (Andy) Gan · GitHub · 901e0c24 · d2e1cfce
Unverified Commit d2e1cfce authored May 27, 2020 by Quan (Andy) Gan Committed by GitHub May 27, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

src/array/array.cc src/array/array.cc +2 -2

No files found.
--- a/src/array/array.cc
+++ b/src/array/array.cc
@@ -613,8 +613,8 @@ COOMatrix COORowWiseTopk(
    COOMatrix mat, IdArray rows, int64_t k, FloatArray weight, bool ascending) {
  COOMatrix ret;
  ATEN_COO_SWITCH(mat, XPU, IdType, {
-    ATEN_FLOAT_TYPE_SWITCH(weight->dtype, FloatType, "weight", {
-      ret = impl::COORowWiseTopk<XPU, IdType, FloatType>(
+    ATEN_DTYPE_SWITCH(weight->dtype, DType, "weight", {
+      ret = impl::COORowWiseTopk<XPU, IdType, DType>(
          mat, rows, k, weight, ascending);
    });
  });