official/mnist: Add input_shape parameter to initial layer (#4104) (#4118)

04ef5c2f · eantaev · Katherine Wu · 30580a35 · 04ef5c2f
Commit 04ef5c2f authored Apr 30, 2018 by eantaev Committed by Katherine Wu Apr 30, 2018
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 1 deletion

official/mnist/mnist.py official/mnist/mnist.py +3 -1

No files found.
--- a/official/mnist/mnist.py
+++ b/official/mnist/mnist.py
@@ -62,7 +62,9 @@ def create_model(data_format):
  # (a subclass of tf.keras.Model) makes for a compact description.
  return tf.keras.Sequential(
      [
-          l.Reshape(input_shape),
+          l.Reshape(
+              target_shape=input_shape,
+              input_shape=(28 * 28,)),
          l.Conv2D(
              32,
              5,