code: def get_model_fn(nn, learningRate, num_classes=None):
      """Sets up loss and optimization details for training and evaluating tensorflow neural network."""
  
      def my_model_fn(features, labels, mode, nnet, lr, num_classes):
          # Getting tf.estimator.EstimatorSpec for PREDICT, TRAIN, and EVAL modes
          op_layer = mlp_definition(features, nn, num_classes)
          # following predictions are used for multiple modes (PREDICT and EVAL)
          predictions = {
              "classes": tf.argmax(input=op_layer, axis=1),
              "probabilities": op_layer,
          }
          label_predictions = {
              "classes": tf.argmax(input=labels, axis=1)
          }
          # Make sure labels are correct type
          if nn.class_or_reg == "reg":
              labels = tf.cast(labels, tf.float32)
          elif nn.class_or_reg == "class":
              labels = tf.cast(labels, tf.int32)
          # For mode = PREDICT
          if mode == tf.estimator.ModeKeys.PREDICT:
              return tf.estimator.EstimatorSpec(
                  mode=mode, predictions=predictions
              )
          # Calculate Loss (for both mode = TRAIN and EVAL)
          if nn.class_or_reg == "reg":
              loss = tf.losses.absolute_difference(
                  labels=labels, predictions=op_layer
              )  # Specify loss here
          elif nn.class_or_reg == "class":
              loss = tf.losses.softmax_cross_entropy(
                  onehot_labels=labels, logits=op_layer
              )  # Specify loss here
              # loss = tf.losses.absolute_difference(labels=labels, predictions=op_layer) # Specify loss here
              # loss = tf.losses.sigmoid_cross_entropy(multi_class_labels=labels, logits=op_layer) # Specify loss here
          # For mode = TRAIN
          if mode == tf.estimator.ModeKeys.TRAIN:
              # Set learning rate
              optimizer = tf.train.GradientDescentOptimizer(
                  learning_rate=learningRate
              )
              train_op = optimizer.minimize(
                  loss=loss, global_step=tf.train.get_global_step()
              )
              return tf.estimator.EstimatorSpec(
                  mode=mode, loss=loss, train_op=train_op
              )
          if nn.class_or_reg == "reg":
              eval_metric_ops = {
                  "mse": tf.metrics.mean_squared_error(
                      labels=labels,
                      predictions=predictions["probabilities"],
                  )
              }
          elif nn.class_or_reg == "class":
              labels = tf.cast(labels, tf.int32)
              eval_metric_ops = {
                  "accuracy": tf.metrics.accuracy(
                      labels=label_predictions["classes"],
                      predictions=predictions["classes"],
                  )
              }
          return tf.estimator.EstimatorSpec(
              mode=mode, loss=loss, eval_metric_ops=eval_metric_ops
          )
  
      def _get_my_model_fn(_nn, _learningRate, _num_classes):
          return lambda features, labels, mode: my_model_fn(
              features, labels, mode, _nn, _learningRate, _num_classes
          )
  
      ret = _get_my_model_fn(nn, learningRate, num_classes)
      return ret
  
file_path: dragonfly/examples/nas/cg/run_tensorflow_testset.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  import tensorflow
  import numpy
  import os
label_desc: Use image node_id 3 for calling a global function with example usage: get_model_fn(nn, learningRate, num_classes) and returns: ret, tf, function, tf, tf