diff --git a/tests/utils.py b/tests/utils.py
index 8bdf55af..7465b396 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -159,7 +159,7 @@ def run_prediction(dataloader, trained_model, dp=False, min_acc=0.50):
         acc = torch.mean(acc).item()
 
     else:
-        y_hat = trained_model(x, 0)
+        y_hat = trained_model(x)
 
         # acc
         labels_hat = torch.argmax(y_hat, dim=1)