diff --git a/examples/mnist/keras/mnist_mlp_estimator.py b/examples/mnist/keras/mnist_mlp_estimator.py
index 7f3329b6..570dad4b 100644
--- a/examples/mnist/keras/mnist_mlp_estimator.py
+++ b/examples/mnist/keras/mnist_mlp_estimator.py
@@ -104,6 +104,22 @@ def predict_input_fn():
       for result in predictions:
         tf_feed.batch_results([result])
 
+  # WORKAROUND FOR https://github.com/tensorflow/tensorflow/issues/21745
+  # wait for all other nodes to complete (via done files)
+  done_dir = "{}/{}/done".format(ctx.absolute_path(args.model), args.mode)
+  print("Writing done file to: {}".format(done_dir))
+  tf.gfile.MakeDirs(done_dir)
+  with tf.gfile.GFile("{}/{}".format(done_dir, ctx.task_index), 'w') as done_file:
+    done_file.write("done")
+
+  for i in range(60):
+    if len(tf.gfile.ListDirectory(done_dir)) < len(ctx.cluster_spec['worker']):
+      print("{} Waiting for other nodes {}".format(datetime.now().isoformat(), i))
+      time.sleep(1)
+    else:
+      print("{} All nodes done".format(datetime.now().isoformat()))
+      break
+
 
 if __name__ == '__main__':
   import argparse
diff --git a/pom.xml b/pom.xml
index 5bcf35e8..99f3d86e 100644
--- a/pom.xml
+++ b/pom.xml
@@ -5,7 +5,7 @@
     <modelVersion>4.0.0</modelVersion>
     <groupId>com.yahoo.ml</groupId>
     <artifactId>tensorflowonspark</artifactId>
-    <version>1.0</version>
+    <version>1.0.1</version>
     <packaging>jar</packaging>
     <name>tensorflowonspark</name>
     <description>Spark Scala inferencing for TensorFlowOnSpark</description>
@@ -25,7 +25,7 @@
         <json4s-native.version>3.5.3</json4s-native.version>
         <maven-shade-plugin.version>2.1</maven-shade-plugin.version>
         <maven-surefire-plugin.version>2.20.1</maven-surefire-plugin.version>
-        <spark.version>2.1.0</spark.version>
+        <spark.version>[2.2.0,)</spark.version>
         <scala.version>2.11.8</scala.version>
         <scala-maven-plugin.version>3.2.1</scala-maven-plugin.version>
         <scala-parser-combinators.version>1.1.0</scala-parser-combinators.version>