message: Always add layer annotations, regardless of mode.
diff: @@ -24,7 +24,6 @@ import pickle
  from google.protobuf.any_pb2 import Any
  from tensorflow.python.estimator import estimator
  -from tensorflow.python.estimator import model_fn
  from tensorflow.python.estimator.canned import dnn
  from tensorflow.python.feature_column import feature_column as feature_column_lib
  from tensorflow.python.framework import ops
  @@ -68,7 +67,7 @@ def _to_any_wrapped_tensor_info(tensor):
  return any_buf
  -def make_input_layer_with_layer_annotations(original_input_layer, mode):
  +def make_input_layer_with_layer_annotations(original_input_layer):
  """Make an input_layer replacement function that adds layer annotations."""
  def input_layer_with_layer_annotations(features,
  @@ -137,9 +136,6 @@ def make_input_layer_with_layer_annotations(original_input_layer, mode):
  if cols_to_output_tensors is not None:
  cols_to_output_tensors = local_cols_to_output_tensors
  - if mode and mode == model_fn.ModeKeys.PREDICT:
  - # Only annotate in PREDICT mode.
  -
  # Annotate features.
  # These are the parsed Tensors, before embedding.
  @@ -148,7 +144,7 @@ def make_input_layer_with_layer_annotations(original_input_layer, mode):
  # spec and looking at the keys.
  spec = feature_column_lib.make_parse_example_spec(feature_columns)
  for key in spec.keys():
  - tensor = features[key]
  + tensor = ops.convert_to_tensor(features[key])
  ops.add_to_collection(
  LayerAnnotationsCollectionNames.keys(
  LayerAnnotationsCollectionNames.UNPROCESSED_FEATURES), key)
  @@ -167,8 +163,7 @@ def make_input_layer_with_layer_annotations(original_input_layer, mode):
  for column, tensor in local_cols_to_output_tensors.items():
  ops.add_to_collection(
  LayerAnnotationsCollectionNames.keys(
  - LayerAnnotationsCollectionNames.PROCESSED_FEATURES),
  - column.name)
  + LayerAnnotationsCollectionNames.PROCESSED_FEATURES), column.name)
  ops.add_to_collection(
  LayerAnnotationsCollectionNames.values(
  LayerAnnotationsCollectionNames.PROCESSED_FEATURES),
  @@ -302,8 +297,8 @@ def DNNClassifierWithLayerAnnotations( # pylint: disable=invalid-name
  def _model_fn(features, labels, mode, config):
  with _monkey_patch(
  feature_column_lib, 'input_layer',
  - make_input_layer_with_layer_annotations(feature_column_lib.input_layer,
  - mode)):
  + make_input_layer_with_layer_annotations(
  + feature_column_lib.input_layer)):
  return original.model_fn(features, labels, mode, config)
  return estimator.Estimator(
  @@ -423,8 +418,8 @@ def DNNRegressorWithLayerAnnotations( # pylint: disable=invalid-name
  def _model_fn(features, labels, mode, config):
  with _monkey_patch(
  feature_column_lib, 'input_layer',
  - make_input_layer_with_layer_annotations(feature_column_lib.input_layer,
  - mode)):
  + make_input_layer_with_layer_annotations(
  + feature_column_lib.input_layer)):
  return original.model_fn(features, labels, mode, config)
  return estimator.Estimator(
  