Skip to content

GitLab

Commit 5ec4c759 authored by Neal Wu's avatar Neal Wu
Browse files

Fix mnist_tpu.py which was giving unknown shape errors

parent 40d29d42
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 0 deletions
official/mnist/dataset.py
View file @ 5ec4c759
...@@ -91,6 +91,7 @@ def dataset(directory, images_file, labels_file): ...@@ -91,6 +91,7 @@ def dataset(directory, images_file, labels_file):
def decode_label(label): def decode_label(label):
label = tf.decode_raw(label, tf.uint8) # tf.string -> [tf.uint8] label = tf.decode_raw(label, tf.uint8) # tf.string -> [tf.uint8]
label = tf.reshape(label, []) # label is a scalar
return tf.to_int32(label) return tf.to_int32(label)
images = tf.data.FixedLengthRecordDataset( images = tf.data.FixedLengthRecordDataset(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment