Iterating dataset as numpy with batch_size > 1 failed to work.

Short description When I loaded a dataset (cifar10) with batch_size set to larger than 1, converted it to a generator with tfds.as_numpy and tried to iterate over the generator, I got the following repeated warning and no data:

2019-03-26 10:09:57.170947: W tensorflow/core/framework/model.cc:419] Failed to find a tunable parameter that would decrease the output time, aborting the current optimization attempt.

Environment information

• Operating System: Mac OSX
• Python version: 2.7
• tensorflow-datasets/tfds-nightly version: 1.0.1
• tensorflow/tensorflow-gpu/tf-nightly/tf-nightly-gpu version: 1.13.1

Reproduction instructions

import tensorflow_datasets as tfds

BATCH_SIZE = 5
dataset = tfds.load('cifar10', as_supervised=True, split=tfds.Split.TEST, batch_size=BATCH_SIZE)
npds = tfds.as_numpy(dataset)

for image, label in npds:
    print(image.shape, label)

Expected behavior The generator npds should return batches of data with the same batch size as dataset.

[pat-coady]

Try this instead:

import tensorflow_datasets as tfds

BATCH_SIZE = 5
dataset = tfds.load('cifar10', as_supervised=True, split=tfds.Split.TEST)
dataset = dataset.batch(BATCH_SIZE)
npds = tfds.as_numpy(dataset)

for image, label in npds:
    print(image.shape, label)

I had another strange problem with the batch_size argument. The latest update made my old problem go away, but then I saw the same problem as you.

[kyleabeauchamp]

I’m using python 3.7 and tf2.0 nightly