Unable to use create_from_images_raw in training

[aiXander]

Minimal code to reproduce this error:

import tensorflow as tf
import cv2
import numpy as np

shape = (3, 1024, 1024)
img = np.random.random((1024, 1024,3)) * 255
img = img.astype(np.uint8)
cv2.imwrite('test.jpg', img)

with tf.gfile.FastGFile('test.jpg', 'rb') as fid:
    encoded_jpg = fid.read()

ex = tf.train.Example(
    features=tf.train.Features(
        feature={
            "img":tf.train.Feature(bytes_list=tf.train.BytesList(value=[encoded_jpg]))
        }
    )
)

tfr_opt = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.NONE)
tfr_writer = tf.python_io.TFRecordWriter("test.tfrecords", tfr_opt)
tfr_writer.write(ex.SerializeToString())
tfr_writer.close()

for record in tf.python_io.tf_record_iterator("test.tfrecords", tfr_opt):
    ex = tf.train.Example()
    ex.ParseFromString(record)
    data = ex.features.feature["img"].bytes_list.value[0]

    img = np.fromstring(data, np.uint8).reshape(shape)