diff --git a/tf_adapter/python/npu_bridge/estimator/npu/npu_optimizer.py b/tf_adapter/python/npu_bridge/estimator/npu/npu_optimizer.py
index 2598b6f5f082f70ef597508d066ab4d6a4f955cd..fa6b784fbfeffa0c5b39dd771294b0bc107928cd 100644
--- a/tf_adapter/python/npu_bridge/estimator/npu/npu_optimizer.py
+++ b/tf_adapter/python/npu_bridge/estimator/npu/npu_optimizer.py
@@ -10,6 +10,7 @@ import os
 import tensorflow as tf
 from tensorflow.python.eager import context
 from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops.cond_v2 import cond_v2
 from tensorflow.python.ops import gen_control_flow_ops
 from tensorflow.python.ops import gen_math_ops
 from tensorflow.python.ops import math_ops
@@ -211,9 +212,9 @@ class NPUOptimizer(optimizer.Optimizer):
 
         return true_apply_gradients(grads_and_vars, global_step, name)
 
-      update_vars = control_flow_ops.cond(self._is_overall_finite,
-                                          true_apply_gradients_fn,
-                                          gen_control_flow_ops.no_op)
+      update_vars = cond_v2(self._is_overall_finite,
+                            true_apply_gradients_fn,
+                            gen_control_flow_ops.no_op)
 
       # Potentially adjust gradient scale in case of finite gradients.
       return control_flow_ops.group(