From 1be6bec45f025f30aefc4a3460d5052e54df9ab0 Mon Sep 17 00:00:00 2001 From: panghongjun Date: Fri, 5 Feb 2021 17:08:34 +0800 Subject: [PATCH] cond change to cond_v2 --- .../python/npu_bridge/estimator/npu/npu_optimizer.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) 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 2598b6f5f..fa6b784fb 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( -- Gitee