diff --git a/posteriors/laplace/diag_ggn.py b/posteriors/laplace/diag_ggn.py
index 26f5a116..36067a16 100644
--- 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,