From 4868ddd508a567a497935378956e9da18976f152 Mon Sep 17 00:00:00 2001 From: Rohan Jain Date: Thu, 19 Apr 2018 15:32:37 -0700 Subject: Simplifying cols_to_vars update PiperOrigin-RevId: 193585237 --- tensorflow/python/feature_column/feature_column.py | 6 ++---- tensorflow/python/feature_column/feature_column_test.py | 6 ++---- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/tensorflow/python/feature_column/feature_column.py b/tensorflow/python/feature_column/feature_column.py index 87a52f8441..a7c4eabcb2 100644 --- a/tensorflow/python/feature_column/feature_column.py +++ b/tensorflow/python/feature_column/feature_column.py @@ -417,10 +417,8 @@ def linear_model(features, trainable=trainable, name='linear_model') retval = linear_model_layer(features) # pylint: disable=not-callable - if cols_to_vars is None: - return retval - for k, v in linear_model_layer.cols_to_vars().items(): - cols_to_vars[k] = v + if cols_to_vars is not None: + cols_to_vars.update(linear_model_layer.cols_to_vars()) return retval diff --git a/tensorflow/python/feature_column/feature_column_test.py b/tensorflow/python/feature_column/feature_column_test.py index 49e06b8245..d963dd9b55 100644 --- a/tensorflow/python/feature_column/feature_column_test.py +++ b/tensorflow/python/feature_column/feature_column_test.py @@ -1269,10 +1269,8 @@ def get_keras_linear_model_predictions(features, trainable, name='linear_model') retval = keras_linear_model(features) # pylint: disable=not-callable - if cols_to_vars is None: - return retval - for k, v in keras_linear_model.cols_to_vars().items(): - cols_to_vars[k] = v + if cols_to_vars is not None: + cols_to_vars.update(keras_linear_model.cols_to_vars()) return retval -- cgit v1.2.3