Fix reduction with a specified axis.

PiperOrigin-RevId: 286185937 Change-Id: I56bed77a84b758f204888102e8a7d7175b6e3d9e
2019-12-18 07:16:34 -08:00 · 2019-12-18 07:16:34 -08:00 · 949b18f0fe
commit 949b18f0fe
parent 6da1b29ae3
1 changed files with 1 additions and 1 deletions
--- a/tensorflow/python/distribute/distribute_lib.py
+++ b/tensorflow/python/distribute/distribute_lib.py
@ -840,7 +840,7 @@ class Strategy(object):
        if dim is not None:
          # By returning a python value in the static shape case, we can
          # maybe get a fast path for reducing the denominator.
-          return numer, dim
+          return numer, array_ops.constant(dim, dtype=dtypes.int64)
      elif axis < 0:
        axis = axis + array_ops.rank(v)
      if v.shape.rank == 1: