code: def main(layer1, layer2):
      # args = parser.parse_args(argv[1:])
  
      # Fetch the data
      (train_x, train_y), (test_x, test_y) = iris_data.load_data()
  
      # Feature columns describe how to use the input.
      my_feature_columns = []
      for key in train_x.keys():
          my_feature_columns.append(
              tf.feature_column.numeric_column(key=key)
          )
  
      # Build 2 hidden layer DNN with 10, 10 units respectively.
      classifier = tf.estimator.DNNClassifier(
          feature_columns=my_feature_columns,
          # Two hidden layers of 10 nodes each.
          hidden_units=[layer1, layer2],
          # The model must choose between 3 classes.
          n_classes=3,
      )
  
      # Train the Model.
      classifier.train(
          input_fn=lambda: iris_data.train_input_fn(
              train_x, train_y, batch_size
          ),
          steps=train_steps,
      )
  
      # Evaluate the model.
      eval_result = classifier.evaluate(
          input_fn=lambda: iris_data.eval_input_fn(
              test_x, test_y, batch_size
          )
      )
  
      print(
          "\nTest set accuracy: {accuracy:0.3f}\n".format(**eval_result)
      )
  
      # Generate predictions from the model
      expected = ["Setosa", "Versicolor", "Virginica"]
      predict_x = {
          "SepalLength": [5.1, 5.9, 6.9],
          "SepalWidth": [3.3, 3.0, 3.1],
          "PetalLength": [1.7, 4.2, 5.4],
          "PetalWidth": [0.5, 1.5, 2.1],
      }
  
      predictions = classifier.predict(
          input_fn=lambda: iris_data.eval_input_fn(
              predict_x, labels=None, batch_size=batch_size
          )
      )
  
      template = '\nPrediction is "{}" ({:.1f}%), expected "{}"'
  
      for pred_dict, expec in zip(predictions, expected):
          class_id = pred_dict["class_ids"][0]
          probability = pred_dict["probabilities"][class_id]
  
          print(
              template.format(
                  iris_data.SPECIES[class_id], 100 * probability, expec
              )
          )
      return probability
  
file_path: auptimizer/Examples/tf_iris_diff_opt/premade_estimator_wrapper.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import argparse
  import tensorflow
  import iris_data
label_desc: Use image node_id 1 for calling a global function with example usage: main(layer1, layer2) and returns: probability