From 3246352eacff22d5d443c19d1c656dd07e153664 Mon Sep 17 00:00:00 2001 From: Matthias Fey Date: Tue, 12 Sep 2023 06:52:54 +0200 Subject: [PATCH] Fix shape of `node_weight` in `METIS` (#342) * update * update --- .github/workflows/testing.yml | 1 + torch_sparse/metis.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index ac84a64e..8e412cb4 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -48,6 +48,7 @@ jobs: - name: Install main package run: | + pip install scipy python setup.py develop env: WITH_METIS: 1 diff --git a/torch_sparse/metis.py b/torch_sparse/metis.py index 1876a418..6757bf71 100644 --- a/torch_sparse/metis.py +++ b/torch_sparse/metis.py @@ -52,7 +52,7 @@ def partition( value = None if balance_edge: - node_weight = col.new_zeros(col.size(0)) + node_weight = col.new_zeros(rowptr.numel() - 1) node_weight.scatter_add_(0, col, torch.ones_like(col)) if node_weight is not None: