diff --git a/sdk/python/kubeflow/training/api/training_client.py b/sdk/python/kubeflow/training/api/training_client.py
index 8750fcbb2d..15548a23b3 100644
--- a/sdk/python/kubeflow/training/api/training_client.py
+++ b/sdk/python/kubeflow/training/api/training_client.py
@@ -904,9 +904,9 @@ def wait_for_job_conditions(
             if callback:
                 callback(job)
 
-            # Raise an exception if Job is Failed and Failed is not expected condition.
+            # Raise an exception if Job is Failed and Failed is not the expected condition.
             if (
-                constants.JOB_CONDITION_FAILED not in conditions
+                constants.JOB_CONDITION_FAILED not in expected_conditions
                 and utils.has_condition(conditions, constants.JOB_CONDITION_FAILED)
             ):
                 raise RuntimeError(