Merge pull request #21716 from artsobolev:pfor-softplus-fix

PiperOrigin-RevId: 211700362

1 files changed, 2 insertions, 2 deletions

diff --git a/tensorflow/python/ops/parallel_for/pfor.py b/tensorflow/python/ops/parallel_for/pfor.py
index 3c914f6ff6..f9153b6d7d 100644
--- a/tensorflow/python/ops/parallel_for/pfor.py
+++ b/tensorflow/python/ops/parallel_for/pfor.py
@@ -21,8 +21,6 @@ from __future__ import print_function
 
 import collections
 
-from absl import flags
-
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
@@ -41,6 +39,7 @@ from tensorflow.python.ops import nn_ops
 from tensorflow.python.ops import parsing_ops
 from tensorflow.python.ops import sparse_ops
 from tensorflow.python.ops import tensor_array_ops
+from tensorflow.python.platform import flags
 from tensorflow.python.platform import tf_logging as logging
 from tensorflow.python.util import nest
 
@@ -2013,6 +2012,7 @@ def _convert_biasaddgrad(pfor_input):
 @RegisterPForWithArgs("ReluGrad")
 @RegisterPForWithArgs("TanhGrad")
 @RegisterPForWithArgs("SigmoidGrad")
+@RegisterPForWithArgs("SoftplusGrad")
 def _convert_grads(pfor_input, op_type, *args, **kw_args):
   del args
   del kw_args