1 files changed, 22 insertions, 15 deletions
diff --git a/tensorflow/examples/skflow/iris_run_config.py b/tensorflow/examples/skflow/iris_run_config.py
index dff0daf9e8..c678c7c738 100644
--- a/tensorflow/examples/skflow/iris_run_config.py
+++ b/tensorflow/examples/skflow/iris_run_config.py
@@ -16,24 +16,31 @@ from __future__ import division
 from __future__ import print_function
 
 from sklearn import datasets, metrics, cross_validation
+import tensorflow as tf
 
-from tensorflow.contrib import learn
 
+def main(unused_argv):
+  # Load dataset.
+  iris = datasets.load_iris()
+  x_train, x_test, y_train, y_test = cross_validation.train_test_split(
+      iris.data, iris.target, test_size=0.2, random_state=42)
 
-# Load dataset.
-iris = datasets.load_iris()
-X_train, X_test, y_train, y_test = cross_validation.train_test_split(iris.data, iris.target,
-    test_size=0.2, random_state=42)
+  # You can define you configurations by providing a RunConfig object to
+  # estimator to control session configurations, e.g. num_cores
+  # and gpu_memory_fraction
+  run_config = tf.contrib.learn.estimators.RunConfig(
+      num_cores=3, gpu_memory_fraction=0.6)
 
-# You can define you configurations by providing a RunConfig object to
-# estimator to control session configurations, e.g. num_cores and gpu_memory_fraction
-run_config = learn.estimators.RunConfig(num_cores=3, gpu_memory_fraction=0.6)
+  # Build 3 layer DNN with 10, 20, 10 units respectively.
+  classifier = tf.contrib.learn.DNNClassifier(hidden_units=[10, 20, 10],
+                                              n_classes=3,
+                                              config=run_config)
 
-# Build 3 layer DNN with 10, 20, 10 units respectively.
-classifier = learn.TensorFlowDNNClassifier(hidden_units=[10, 20, 10],
-    n_classes=3, steps=200, config=run_config)
+  # Fit and predict.
+  classifier.fit(x_train, y_train, steps=200)
+  score = metrics.accuracy_score(y_test, classifier.predict(x_test))
+  print('Accuracy: {0:f}'.format(score))
 
-# Fit and predict.
-classifier.fit(X_train, y_train)
-score = metrics.accuracy_score(y_test, classifier.predict(X_test))
-print('Accuracy: {0:f}'.format(score))
+
+if __name__ == '__main__':
+  tf.app.run()