Replace partial with a lambda

implying we relax from 'batch' named kwarg to any arg for the batch in 'forward'
normal-computing · May 24, 2024 · 3e8036d · 3e8036d
1 parent 3ab0b67
commit 3e8036d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/posteriors/laplace/diag_ggn.py b/posteriors/laplace/diag_ggn.py
@@ -134,8 +134,8 @@ def outer_loss(z, batch):
 
     with torch.no_grad(), CatchAuxError():
         diag_ggn_batch, aux = diag_ggn(
-            partial(forward, batch=batch),
-            partial(outer_loss, batch=batch),
+            lambda params: forward(params, batch),
+            lambda z: outer_loss(z, batch),
             forward_has_aux=True,
             loss_has_aux=False,
             normalize=False,