diff --git a/keras_estimator_vgg16-cat_vs_dog-TFRecord.ipynb b/keras_estimator_vgg16-cat_vs_dog-TFRecord.ipynb
index 1e6a986..5bc90d0 100644
--- a/keras_estimator_vgg16-cat_vs_dog-TFRecord.ipynb
+++ b/keras_estimator_vgg16-cat_vs_dog-TFRecord.ipynb
@@ -575,7 +575,7 @@
     "        image = tf.reshape(image, image_shape)\n",
     "        image = tf.subtract(image, 116.779) # Zero-center by mean pixel\n",
     "        image = tf.reverse(image, axis=[2]) # 'RGB'->'BGR'\n",
-    "        d = dict(zip([input_name], [image])), [label]\n",
+    "        d={input_name:image},label\n",
     "        return d\n",
     "    \n",
     "    dataset = tf.data.TFRecordDataset(filenames=filenames)\n",
@@ -628,7 +628,7 @@
     "next_batch = imgs_input_fn(path_tfrecords_test, perform_shuffle=True, batch_size=20)\n",
     "with tf.Session() as sess:\n",
     "    first_batch = sess.run(next_batch)\n",
-    "x_d = first_batch[0]['input_1']\n",
+    "x_d = first_batch[0][input_name]\n",
     "\n",
     "print(x_d.shape)\n",
     "img = image.array_to_img(x_d[8])\n",