diff options
-rw-r--r-- | tensorflow/python/feature_column/feature_column.py | 6 | ||||
-rw-r--r-- | 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 |