I am trying to read a TFRecord file like this:
dataset = tf.data.TFRecordDataset("./tfrecords/train.record").map(_extract_fn).batch(3)

However, when I run
features, labels = iter(dataset).next()
I get this error:
tensorflow.python.framework.errors_impl.InvalidArgumentError: Cannot add tensor to the batch: number of elements does not match. Shapes are: [tensor]: [4], [batch]: [2] [Op:IteratorGetNext]

This is the function that parses the TFRecord file:
    features = {
        'image/height': tf.io.FixedLenFeature([],tf.int64),
        'image/width': tf.io.FixedLenFeature([], tf.int64),
        'image/filename': tf.io.VarLenFeature(tf.string),
        'image/id': tf.io.FixedLenFeature([], tf.string),
        'image/encoded': tf.io.FixedLenFeature([],tf.string),
        'image/format': tf.io.FixedLenFeature([], tf.string),
        'image/object/class/text': tf.io.VarLenFeature(tf.string),
        'image/object/class/label': tf.io.VarLenFeature(tf.int64),
        'image/object/bbox/xmin': tf.io.VarLenFeature(tf.float32),
        'image/object/bbox/xmax': tf.io.VarLenFeature(tf.float32),
        'image/object/bbox/ymin': tf.io.VarLenFeature(tf.float32),
        'image/object/bbox/ymax': tf.io.VarLenFeature(tf.float32),
        }
sample = tf.io.parse_single_example(tfrecord, features)
data = {}


data["image/encoded"] = tf.image.decode_jpeg(sample["image/encoded"], channels=3)
label = sample['image/object/class/label'].values

return data,label

If I write return data instead and only set features = iter(dataset).next() it works fine.
What is the issue here?
Thanks for any help!