diff --git a/tests/unit/examples/test_scaling_criteo_merlin_models.py b/tests/unit/examples/test_scaling_criteo_merlin_models.py
index 23e6c495c..5ff786e8b 100644
--- a/tests/unit/examples/test_scaling_criteo_merlin_models.py
+++ b/tests/unit/examples/test_scaling_criteo_merlin_models.py
@@ -59,6 +59,7 @@ def test_func():
                 "precision",
                 "recall",
                 "regularization_loss",
+                "loss_batch",
             ]
         )
         assert os.path.isfile("/tmp/output/criteo/dlrm/saved_model.pb")