Update test to capture the new desired behaviour.

normal-computing · May 24, 2024 · 3ab0b67 · 3ab0b67
1 parent f327548
commit 3ab0b67
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/laplace/test_diag_ggn.py b/tests/laplace/test_diag_ggn.py
@@ -18,8 +18,8 @@ def normal_log_likelihood(y_pred, batch):
     )  # validate args introduces control flows not yet supported in torch.func.vmap
 
 
-def forward_m(params, batch, model):
-    y_pred = functional_call(model, params, batch[0])
+def forward_m(params, b, model):
+    y_pred = functional_call(model, params, b[0])
     return y_pred, torch.tensor([])