From 91e8334601d7ae03a12fecc1a0cc49878a576d9e Mon Sep 17 00:00:00 2001
From: Khue Le <huudienkhue.le@gmail.com>
Date: Thu, 18 Apr 2024 19:04:29 +0200
Subject: [PATCH] Add support for `.eval()`

---
 submission_runner.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/submission_runner.py b/submission_runner.py
index a6f8c05a3..fb63b5a32 100644
--- a/submission_runner.py
+++ b/submission_runner.py
@@ -383,6 +383,8 @@ def train_once(
 
         try:
           eval_start_time = get_time()
+          if hasattr(optimizer_state['optimizer'], 'eval'):
+            optimizer_state['optimizer'].eval()
           latest_eval_result = workload.eval_model(global_eval_batch_size,
                                                    model_params,
                                                    model_state,