optional prepare_for_eval arg

mlcommons · Oct 21, 2024 · 4d74d2c · 4d74d2c
1 parent 9caedc5
commit 4d74d2c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/submission_runner.py b/submission_runner.py
@@ -200,7 +200,7 @@ def train_once(
     init_optimizer_state: spec.InitOptimizerFn,
     update_params: spec.UpdateParamsFn,
     data_selection: spec.DataSelectionFn,
-    prepare_for_eval: spec.PrepareForEvalFn,
+    prepare_for_eval: Optional[spec.PrepareForEvalFn],
     hyperparameters: Optional[spec.Hyperparameters],
     rng_seed: int,
     rng: spec.RandomState,