code: def model_fn(features, labels, mode, config):
      """Model function for custom estimator."""
      del labels
      del config
      classes = features["classes"]
      scores = features["scores"]
  
      with tf.control_dependencies(
          [tf.assert_less(tf.shape(classes)[0], tf.constant(2))]
      ):
          scores = tf.identity(scores)
  
      predictions = {
          prediction_keys.PredictionKeys.LOGITS: scores,
          prediction_keys.PredictionKeys.PROBABILITIES: scores,
          prediction_keys.PredictionKeys.PREDICTIONS: scores,
          prediction_keys.PredictionKeys.CLASSES: classes,
      }
  
      if mode == tf_estimator.ModeKeys.PREDICT:
          return tf_estimator.EstimatorSpec(
              mode=mode,
              predictions=predictions,
              export_outputs={
                  tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY: tf_estimator.export.ClassificationOutput(
                      scores=scores, classes=classes
                  ),
              },
          )
  
      loss = tf.constant(0.0)
      train_op = tf.compat.v1.assign_add(
          tf.compat.v1.train.get_global_step(), 1
      )
      eval_metric_ops = {
          metric_keys.MetricKeys.LOSS_MEAN: tf.compat.v1.metrics.mean(
              loss
          ),
      }
  
      return tf_estimator.EstimatorSpec(
          mode=mode,
          loss=loss,
          train_op=train_op,
          predictions=predictions,
          eval_metric_ops=eval_metric_ops,
      )
  
file_path: model-analysis/tensorflow_model_analysis/eval_saved_model/example_trainers/batch_size_limited_classifier.py
imports: import tensorflow
  from tensorflow import estimator
  from tensorflow_model_analysis.eval_saved_model import export
  from tensorflow_model_analysis.eval_saved_model.example_trainers import util
  from tensorflow.python.estimator.canned import metric_keys
  from tensorflow.python.estimator.canned import prediction_keys
label_desc: Use image node_id 1 for calling a global function with example usage: model_fn(features, labels, mode, config) and returns: tf_estimator, tf_estimator