fix the bug of average degree

iMoonLab · Mar 7, 2024 · ac69684 · ac69684
1 parent dc27de3
commit ac69684
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/dhg/nn/convs/hypergraphs/unignn_conv.py b/dhg/nn/convs/hypergraphs/unignn_conv.py
@@ -62,7 +62,9 @@ def forward(self, X: torch.Tensor, hg: Hypergraph) -> torch.Tensor:
         # compute the special degree of hyperedges
         _De = torch.zeros(hg.num_e, device=hg.device)
         # scatter_reduce() is relay on the torch 1.12.1, which may be updated in the future
-        _De = _De.scatter_reduce(0, index=hg.v2e_dst, src=hg.D_v.clone()._values()[hg.v2e_src], reduce="mean")
+        _Dv = hg.D_v.clone()._values()[hg.v2e_src]
+        _De = _De.scatter_reduce(0, index=hg.v2e_dst, src=_Dv, reduce="sum") / _De.scatter_reduce(
+            0, index=hg.v2e_dst, src=(_Dv != 0).type_as(_De), reduce="sum")
         _De = _De.pow(-0.5)
         _De[_De.isinf()] = 1
         Y = _De.view(-1, 1) * Y