Set correct batch_size for metrics in CNN

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)

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:

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(