From 6c6ff740184b6461fac946b3d31eceb6a2dcf50c Mon Sep 17 00:00:00 2001
From: Tiziano <tizianocitro98@gmail.com>
Date: Mon, 2 Mar 2026 13:54:41 +0100
Subject: [PATCH] fix: set correct batch_size for metrics in CNN

---
 hyperbench/hlp/common_neighbors_hlp.py |  5 +++-
 hyperbench/models/hypergcn.py          | 35 +++++++++++++-------------
 hyperbench/nn/conv.py                  |  5 ++--
 3 files changed, 23 insertions(+), 22 deletions(-)

diff --git a/hyperbench/hlp/common_neighbors_hlp.py b/hyperbench/hlp/common_neighbors_hlp.py
index 2c9848a..6889306 100644
--- a/hyperbench/hlp/common_neighbors_hlp.py
+++ b/hyperbench/hlp/common_neighbors_hlp.py
@@ -83,7 +83,10 @@ def __step(self, batch: HData, stage: Stage) -> Tensor:
         """
         scores: Tensor = self.decoder(batch.hyperedge_index, self.node_to_neighbors)
         labels = batch.y
-        batch_size = batch.num_nodes
+
+        # We need to use the number of hyperedges as batch size for logging purposes,
+        # since each hyperedge is a separate prediction
+        batch_size = batch.num_hyperedges
 
         loss = self._compute_loss(scores, labels, batch_size, stage)
         self._compute_metrics(scores, labels, batch_size, stage)
diff --git a/hyperbench/models/hypergcn.py b/hyperbench/models/hypergcn.py
index ee6d300..d421b16 100644
--- a/hyperbench/models/hypergcn.py
+++ b/hyperbench/models/hypergcn.py
@@ -36,24 +36,23 @@ def __init__(
         self.use_mediator = use_mediator
         self.cached_gcn_laplacian_matrix: Optional[Tensor] = None
 
-        self.layers = nn.ModuleList()
-        self.layers.append(
-            HyperGCNConv(
-                in_channels=in_channels,
-                out_channels=hid_channels,
-                use_batch_normalization=use_batch_normalization,
-                drop_rate=drop_rate,
-                use_mediator=use_mediator,
-            )
-        )
-        self.layers.append(
-            HyperGCNConv(
-                in_channels=hid_channels,
-                out_channels=num_classes,
-                use_batch_normalization=use_batch_normalization,
-                use_mediator=use_mediator,
-                is_last=True,
-            )
+        self.layers = nn.ModuleList(
+            [
+                HyperGCNConv(
+                    in_channels=in_channels,
+                    out_channels=hid_channels,
+                    use_batch_normalization=use_batch_normalization,
+                    drop_rate=drop_rate,
+                    use_mediator=use_mediator,
+                ),
+                HyperGCNConv(
+                    in_channels=hid_channels,
+                    out_channels=num_classes,
+                    use_batch_normalization=use_batch_normalization,
+                    use_mediator=use_mediator,
+                    is_last=True,
+                ),
+            ]
         )
 
     def forward(self, x: Tensor, hyperedge_index: Tensor) -> Tensor:
diff --git a/hyperbench/nn/conv.py b/hyperbench/nn/conv.py
index 5d14b74..5a5e984 100644
--- a/hyperbench/nn/conv.py
+++ b/hyperbench/nn/conv.py
@@ -35,9 +35,8 @@ def __init__(
         self.activation_fn = nn.ReLU(inplace=True)
         self.dropout = nn.Dropout(drop_rate)
 
-        # θ is the learnable weight matrix (as in the HyperGCN paper)
-        # # it projects node features from in_channels to out_channels
-        # and learns how to mix feature channels
+        # θ is the learnable weight matrix (as in the HyperGCN paper),
+        # it projects node features from in_channels to out_channels and learns how to mix feature channels
         self.theta = nn.Linear(in_channels, out_channels, bias=bias)
 
     def forward(