keep track of best_val_loss

delphi-suite · Mar 7, 2024 · ae9d8e6 · ae9d8e6
1 parent 6856bed
commit ae9d8e6
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/delphi/train/training.py b/src/delphi/train/training.py
@@ -234,14 +234,14 @@ def train_step(
         )
     iter_num += 1
     local_iter_num += 1
-    return False, t0, iter_num, local_iter_num
+    return False, t0, iter_num, local_iter_num, best_val_loss
 
 
 for epoch in range(config.max_epochs):
     train_ds.shuffle(epoch)
     train_batch_iter = iter(DataLoader(train_ds, batch_size=config.batch_size))  # type: ignore
     for _ in tqdm(range(num_steps)):
-        breaknow, t0, iter_num, local_iter_num = train_step(
+        breaknow, t0, iter_num, local_iter_num, best_val_loss = train_step(
             train_ds,
             validation_ds,
             lr_decay_iters,