From aeb83a33e96908ea61409ac3a2d2856ab498d6e1 Mon Sep 17 00:00:00 2001 From: c00420053 Date: Mon, 8 Jan 2024 21:19:42 +0800 Subject: [PATCH] filter none grad when satmode is inf_nan --- .../npu_bridge/estimator/npu/npu_loss_scale_optimizer.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tf_adapter/python/npu_bridge/estimator/npu/npu_loss_scale_optimizer.py b/tf_adapter/python/npu_bridge/estimator/npu/npu_loss_scale_optimizer.py index ba3435b36..9d9c477a4 100644 --- a/tf_adapter/python/npu_bridge/estimator/npu/npu_loss_scale_optimizer.py +++ b/tf_adapter/python/npu_bridge/estimator/npu/npu_loss_scale_optimizer.py @@ -53,7 +53,8 @@ class NPULossScaleOptimizer(lso.LossScaleOptimizer): def apply_gradients(self, grads_and_vars, global_step=None, name=None): """Apply gradients. See base class `tf.compat.v1.train.Optimizer`.""" if npu_plugin.is_inf_nan_enabled(): - return super().apply_gradients(grads_and_vars, global_step, name) + non_empty_grads_and_vars = [(grad, var) for grad, var in grads_and_vars if grad is not None] + return super().apply_gradients(non_empty_grads_and_vars, global_step, name) if self._enable_overflow_check(): with tf.name_scope(self._name): -- Gitee