From 107c6b6e2ad3312d77ad4e99034f17a31f2967c6 Mon Sep 17 00:00:00 2001
From: Niccolo-Ajroldi <n.ajroldi@gmail.com>
Date: Thu, 3 Oct 2024 18:05:01 +0200
Subject: [PATCH] update test

---
 tests/reference_algorithm_tests.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/reference_algorithm_tests.py b/tests/reference_algorithm_tests.py
index 74c06e180..938c4fa11 100644
--- a/tests/reference_algorithm_tests.py
+++ b/tests/reference_algorithm_tests.py
@@ -471,6 +471,7 @@ def _test_submission(workload_name,
         batch=batch,
         loss_type=workload.loss_type,
         optimizer_state=optimizer_state,
+        train_state={},
         eval_results=[],
         global_step=global_step,
         rng=update_rng)