Add EKFAC tests and fix a couple of bugs

bergson/collector/collector.py

@@ -621,7 +621,7 @@ def fwd_bwd_hessian(model, x: Tensor, y: Tensor, batch: dict):
             if index_cfg.loss_reduction == "mean"
             else 1.0
         )
-        if not hessian_cfg.use_dataset_labels:
+        if hessian_cfg.use_dataset_labels:
             losses = F.cross_entropy(
                 logits.reshape(-1, logits.size(-1)),
                 y[:, 1:].flatten(),

bergson/hessians/kfac.py

@@ -69,9 +69,10 @@ def backward_hook(self, module: nn.Module, g: Tensor) -> None:
         """Compute gradient covariance: G^T @ G."""
         name = assert_type(str, module._name)
         S_cov_po = self.S_cov_dict[name]
+        mask = self._current_valid_mask
 
-        # Reshape to [N*S, O]
-        g_bo = g.reshape(-1, g.shape[-1])
+        # g: [N, S, O], mask: [N, S] -> select valid positions
+        g_bo = g[mask]  # [num_valid, O]
 
         # Compute local covariance
         local_update_oo = g_bo.mT @ g_bo