cleanup

Signed-off-by: Mayank Mishra <[email protected]>
IBM · Jan 8, 2025 · 4845c06 · 4845c06
1 parent 1dfff82
commit 4845c06
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/dolomite_engine/finetune.py b/dolomite_engine/finetune.py
@@ -274,13 +274,15 @@ def evaluate(
 
     metrics_tracker = MetricsTrackingDict({})
     val_dataloader = custom_iterator(val_dataloader, infinite=False)
+    loss_tokens = 0
 
     for _ in range(num_steps):
         batch = get_next_batch(val_dataloader)
+        loss_tokens += (batch["labels"] != -100).sum()
         loss_step_dict = model_container[0](batch)
         metrics_tracker = metrics_tracker + loss_step_dict
 
-    metrics_tracker = metrics_tracker / num_steps
+    metrics_tracker = metrics_tracker / loss_tokens
 
     for key in metrics_tracker:
         metrics_tracker[key] = dtensor_to_tensor(metrics_tracker[key])