The code is:
def _proc_and_batch(self, ds, batch_size):
def _process_data(x_):
img_ = tf.cast(x_['image'], tf.int32)
img_.set_shape(self._img_shape)
return pack(image=img_, label=tf.constant(0, dtype=tf.int32))
ds = ds.map(_process_data, num_parallel_calls=tf.data.experimental.AUTOTUNE)
ds = ds.batch(batch_size, drop_remainder=True)
ds = ds.prefetch(tf.data.experimental.AUTOTUNE)
return ds
How can _process_data be called without giving it an input? So, what's its input?
For reference, the code is taken from here:
https://github.com/hojonathanho/diffusion/blob/master/diffusion_tf/tpu_utils/datasets.py