Internal change.

PiperOrigin-RevId: 292533112
ChanchalKumarMaji · Feb 3, 2020 · 1de053c · 1de053c
1 parent 6484258
commit 1de053c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/mnist_export_v2/export_test.py b/examples/mnist_export_v2/export_test.py
@@ -35,7 +35,7 @@ def setUp(self):
     super(ExportTest, self).setUp()
     def create_image_and_label(index):
       image = tf.cast(255 * tf.random.normal([1, 28, 28, 1]), tf.uint8)
-      return dict(image=image, label=index)
+      return dict(image=image, label=[index])
     self.mock_dataset = tf.data.Dataset.range(5).map(create_image_and_label)
 
   def test_model_exporting(self):