fix random forest TF 1.4 compatibility

jdoe68877 · Dec 12, 2017 · 0c4e666 · 0c4e666
1 parent d3f3c83
commit 0c4e666
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/examples/2_BasicModels/random_forest.py b/examples/2_BasicModels/random_forest.py
@@ -47,15 +47,15 @@
 loss_op = forest_graph.training_loss(X, Y)
 
 # Measure the accuracy
-infer_op = forest_graph.inference_graph(X)
+infer_op, _, _ = forest_graph.inference_graph(X)
 correct_prediction = tf.equal(tf.argmax(infer_op, 1), tf.cast(Y, tf.int64))
 accuracy_op = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
 
 # Initialize the variables (i.e. assign their default value)
 init_vars = tf.global_variables_initializer()
 
 # Start TensorFlow session
-sess = tf.Session()
+sess = tf.train.MonitoredSession()
 
 # Run the initializer
 sess.run(init_vars)

diff --git a/notebooks/2_BasicModels/random_forest.ipynb b/notebooks/2_BasicModels/random_forest.ipynb
@@ -122,7 +122,7 @@
     "loss_op = forest_graph.training_loss(X, Y)\n",
     "\n",
     "# Measure the accuracy\n",
-    "infer_op = forest_graph.inference_graph(X)\n",
+    "infer_op, _, _ = forest_graph.inference_graph(X)\n",
     "correct_prediction = tf.equal(tf.argmax(infer_op, 1), tf.cast(Y, tf.int64))\n",
     "accuracy_op = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))\n",
     "\n",
@@ -158,7 +158,7 @@
    ],
    "source": [
     "# Start TensorFlow session\n",
-    "sess = tf.Session()\n",
+    "sess = tf.train.MonitoredSession()\n",
     "\n",
     "# Run the initializer\n",
     "sess.run(init_vars)\n",