title: How to reshape the (image, label) dataset after py_function
body: <p>I am trying to read a custom mapped dataset for training. But after I map the dataset using a py_function, I get unknown shape like:</p>
  <pre><code>def process_path(file_path):
  
      label = get_label(file_path)
      img = tf.io.read_file(file_path)
      img = decode_img(img)
  
      print('image shape:', img.shape) #this print correctly: image shape: (180, 180, 3)
      print('label shape:', label.shape) #this print correctly: label shape: ()
  
      return img, label
  
  train_ds = train_ds.map(lambda x: tf.py_function(process_path, [x], (tf.float32, tf.int32))) 
  print(train_ds)
  # this print unknown shape &lt;PrefetchDataset shapes: (&lt;unknown&gt;, &lt;unknown&gt;), types: (tf.float32, tf.int32)&gt;
  </code></pre>
  <p>This will fail the model.fit(), so I want to reshape the dataset to the correct shape like:</p>
  <pre><code>&lt;BatchDataset shapes: ((None, 180, 180, 3), (None,)), types: (tf.float32, tf.int32)&gt;
  </code></pre>
  <p>using:</p>
  <pre><code>train_ds = tf.reshape(train_ds, ((None, 180, 180, 3), (None,)))
  </code></pre>
  <p>But this will give an error:</p>
  <pre><code>ValueError: Attempt to convert a value (&lt;MapDataset shapes: (&lt;unknown&gt;, &lt;unknown&gt;), types: (tf.float32, tf.int32)&gt;) with an unsupported type (&lt;class 'tensorflow.python.data.ops.dataset_ops.MapDataset'&gt;) to a Tensor.
  </code></pre>
  <p>How can I correctly assign the (image, label) shape in this step?</p>