From 18a9074060e3e78495e20141effd96c5da732479 Mon Sep 17 00:00:00 2001 From: Trevor Morris Date: Fri, 19 Jul 2019 15:54:40 -0700 Subject: [PATCH] Improve code based on reviewer feedback --- tensorflow/python/keras/layers/core.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tensorflow/python/keras/layers/core.py b/tensorflow/python/keras/layers/core.py index b21801786d9..117e2d9749c 100644 --- a/tensorflow/python/keras/layers/core.py +++ b/tensorflow/python/keras/layers/core.py @@ -580,9 +580,10 @@ class Flatten(Layer): permutation.append(1) inputs = array_ops.transpose(inputs, perm=permutation) - input_shape = tensor_shape.TensorShape(inputs.shape).as_list() - if input_shape and all(input_shape[1:]): - outputs = array_ops.reshape(inputs, (-1, int(np.prod(input_shape[1:])))) + input_shape = inputs.shape + if input_shape[1:].is_fully_defined(): + outputs = array_ops.reshape( + inputs, (-1, tensor_shape.dimension_value(np.prod(input_shape[1:])))) else: outputs = array_ops.reshape( inputs, (tensor_shape.dimension_value(inputs.shape[0]) or