title: Google colab not loading image files while using tensorflow 2.0 batched dataset
body: <p>A little bit of background, I am loading about 60,000 images to colab to train a GAN. I have already uploaded them to Drive and the directory structure contains folders for different classes (about 7-8) inside <code>root</code>. I am loading them to colab as follows:</p>
  
  <pre><code>root = "drive/My Drive/data/images"
  root = pathlib.Path(root)
  
  list_ds = tf.data.Dataset.list_files(str(root/'*/*'))
  
  for f in list_ds.take(3):
    print(f.numpy())
  </code></pre>
  
  <p>which gives the ouput:</p>
  
  <pre><code>b'drive/My Drive/data/images/folder_1/2994.jpg'
  b'drive/My Drive/data/images/folder_1/6628.jpg'
  b'drive/My Drive/data/images/folder_2/37872.jpg'
  </code></pre>
  
  <p>I am further processing them as follows:</p>
  
  <pre><code>def process_path(file_path):
    label = tf.strings.split(file_path, '/')[-2]
    image = tf.io.read_file(file_path)
    image = tf.image.decode_jpeg(image)
    image = tf.image.convert_image_dtype(image, tf.float32)
    return image#, label
  
  ds = list_ds.map(process_path)
  
  BUFFER_SIZE = 60000
  BATCH_SIZE = 128
  
  train_dataset = ds.shuffle(BUFFER_SIZE).batch(BATCH_SIZE)
  </code></pre>
  
  <p>Each image is of size <code>128x128</code>. Now coming to the problem when I try to view a batch in colab the execution goes on forever and never stops, for example, with this code:</p>
  
  <pre><code>for batch in train_dataset.take(4):
    print([arr.numpy() for arr in batch])
  </code></pre>
  
  <p>Earlier I thought that batch_size might be a issue so tried changing it but still same problem. Can it be a problem due to colab as I am loading a large number of files?</p>
  
  <p>Or due to the size of images as it was working when using MNIST(28x28)? If so, what are the possible solutions?</p>
  
  <p>Thanks in advance.</p>
  
  <p>EDIT:
  After removing the shuffle statement, the last line gets executed within a few seconds. So I thought it could be a problem due to BUFFER_SIZE of shuffle, but even with a reduced BUFFER_SIZE, it is again taking a very long time to execute. Any workaround?</p>