title: tf.data How to npy load data for training?
body: <p>I am trying to use tf.data.Dataset on my downloaded dataset made over 1mln audio spectrogram that I have already computed and stored on my HardDisk as npy files. I have stored all the paths on a CSV dataframe and I am trying to use tf.data:</p>
  <pre><code>data_selected=pd.read_csv(data_path)
  filename = data_selected['filename']
  idx_frame=data_selected['index_frame']
  labels_train=data_selected['class']
  </code></pre>
  <p>Then i use from_tensor_slices:</p>
  <pre><code> dataset = tf.data.Dataset.from_tensor_slices((filename, we_objects, idx_frame)).shuffle(len(spect_names))
  </code></pre>
  <p>From it I created a parse function to load:</p>
  <pre><code>def parse_function(filename, label, idx_frame):
  
  Sxx = np.load( filename, allow_pickle = True)
  Sxx = Sxx[idx_frame].reshape(Sxx.shape[1],Sxx.shape[2],1)
  return Sxx, label
  </code></pre>
  <p>But when I do:</p>
  <pre><code>dataset = dataset.map(parse_function, num_parallel_calls= 'auto')
  </code></pre>
  <p>I got error that filename is a tensor and not a string...
  So I have used py_func:</p>
  <pre><code>dataset = dataset.map(lambda item, lab, idx: tf.numpy_function(
        parse_function, [item, lab, idx], [tf.float32, tf.float32, tf.int32],
        num_parallel_calls=tf.data.AUTOTUNE)
  dataset.batch(batch_size)
  dataset = dataset.prefetch(1)
  options = tf.data.Options()
  from tensorflow.data.experimental import AutoShardPolicy
  options.experimental_distribute.auto_shard_policy = AutoShardPolicy.OFF
  dataset = dataset.with_options(options)
  dataset = dataset.cache()
  </code></pre>
  <p>And no errors, but then I do:</p>
  <pre><code> history=model.fit(dataset,
               initial_epoch=i, epochs=end_epoch,
                callbacks=[tbCallBack,model_checkpoint])
  </code></pre>
  <p>And I got errors:</p>
  <pre><code> InvalidArgumentError: 3 root error(s) found.
    (0) Invalid argument:  pyfunc_14 returns 2 values, but expects to see 3 values.
       [[{{node PyFunc}}]]
       [[MultiDeviceIteratorGetNextFromShard]]
       [[RemoteCall]]
       [[IteratorGetNextAsOptional]]
       [[div_no_nan/ReadVariableOp/_270]]
    (1) Invalid argument:  pyfunc_14 returns 2 values, but expects to see 3 values.
       [[{{node PyFunc}}]]
       [[MultiDeviceIteratorGetNextFromShard]]
       [[RemoteCall]]
       [[IteratorGetNextAsOptional]]
       [[replica_1/angular_distance/weighted_loss/cond/switch_pred/_149/_88]]
    (2) Invalid argument:  pyfunc_14 returns 2 values, but expects to see 3 values.
       [[{{node PyFunc}}]]
       [[MultiDeviceIteratorGetNextFromShard]]
       [[RemoteCall]]
       [[IteratorGetNextAsOptional]]
  0 successful operations.
  0 derived errors ignored. [Op:__inference_train_function_51059]
  
  Function call stack:
  train_function -&gt; train_function -&gt; train_function
                                                            
  </code></pre>
  <p>Where is the error? Is it the right way to do?</p>
  <p>Thank you!</p>