Hi,
I’m trying to apply preprocessing to tf.data.Dataset (option 2 in the Keras Preprocessing Guide: Bekerja dengan lapisan preprocessing  |  TensorFlow Core 2). I get an exception thrown when trying to read from a mapped Dataset which includes the RandomFlip preprocessing layer.
Any ideas what am I doing wrong here? code + output below.
Thanks!
import functools
import numpy as np
import tensorflow as tf


def data_gen():
    for i in range(10):
        x = np.random.random(size=(80, 80, 3)) * 255
        y = np.random.random(size=(40, 40, 1)) * 255
        yield x, y


def preprocess(image, label, cropped_image_size, cropped_label_size, skip_augmentations=False):

    x = image
    y = label

    x_size = cropped_image_size
    y_size = cropped_label_size

    if not skip_augmentations:
        x = tf.keras.layers.RandomFlip(mode="horizontal")(x)
        y = tf.keras.layers.RandomFlip(mode="horizontal")(y)

        x = tf.keras.layers.RandomRotation(factor=1.0, fill_mode='constant')(x)
        y = tf.keras.layers.RandomRotation(factor=1.0, fill_mode='constant')(y)

    x = tf.keras.layers.CenterCrop(x_size, x_size)(x)
    y = tf.keras.layers.CenterCrop(y_size, y_size)(y)

    return x, y


print(tf.__version__)
dataset = tf.data.Dataset.from_generator(data_gen, output_signature=(
    tf.TensorSpec(shape=(80, 80, 3), dtype='float32'),
    tf.TensorSpec(shape=(40, 40, 1), dtype='float32')
))

crop_only_fn = functools.partial(preprocess, cropped_image_size=50, cropped_label_size=25, skip_augmentations=True)
train_preprocess_fn = functools.partial(preprocess, cropped_image_size=50, cropped_label_size=25, skip_augmentations=False)

crop_dataset = dataset.map(crop_only_fn)
train_dataset = dataset.map(train_preprocess_fn)

it_crop = iter(crop_dataset)
ex_crop = next(it_crop)         # this works

it_train = iter(train_dataset)
ex_train = next(it_train)       # this fails


Output:
2021-12-13 16:10:15.106947: I tensorflow/core/platform/cpu_feature_guard.cc:151] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2021-12-13 16:10:15.876142: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1525] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 22302 MB memory:  -> device: 0, name: NVIDIA GeForce RTX 3090, pci bus id: 0000:17:00.0, compute capability: 8.6
2021-12-13 16:10:15.877134: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1525] Created device /job:localhost/replica:0/task:0/device:GPU:1 with 22302 MB memory:  -> device: 1, name: NVIDIA GeForce RTX 3090, pci bus id: 0000:65:00.0, compute capability: 8.6
2021-12-13 16:10:16.720083: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at stateful_random_ops.cc:278 : NOT_FOUND: Resource localhost/_AnonymousVar3/N10tensorflow3VarE does not exist.
2021-12-13 16:10:16.720119: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at stateful_random_ops.cc:278 : NOT_FOUND: Resource localhost/_AnonymousVar1/N10tensorflow3VarE does not exist.
2021-12-13 16:10:16.720143: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at stateful_random_ops.cc:278 : NOT_FOUND: Resource localhost/_AnonymousVar0/N10tensorflow3VarE does not exist.
2021-12-13 16:10:16.720199: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at stateful_random_ops.cc:278 : NOT_FOUND: Resource localhost/_AnonymousVar2/N10tensorflow3VarE does not exist.
2.7.0
Traceback (most recent call last):
  File "./issue_dataaug.py", line 51, in <module>
    ex_train = next(it_train)       # this fails
  File "/.../mickey/venvs/cvi/lib/python3.8/site-packages/tensorflow/python/data/ops/iterator_ops.py", line 800, in __next__
    return self._next_internal()
  File "/.../mickey/venvs/cvi/lib/python3.8/site-packages/tensorflow/python/data/ops/iterator_ops.py", line 783, in _next_internal
    ret = gen_dataset_ops.iterator_get_next(
  File "/.../mickey/venvs/cvi/lib/python3.8/site-packages/tensorflow/python/ops/gen_dataset_ops.py", line 2845, in iterator_get_next
    _ops.raise_from_not_ok_status(e, name)
  File "/.../mickey/venvs/cvi/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 7107, in raise_from_not_ok_status
    raise core._status_to_exception(e) from None  # pylint: disable=protected-access
tensorflow.python.framework.errors_impl.NotFoundError: 2 root error(s) found.
  (0) NOT_FOUND: {{function_node __inference_Dataset_map_classfunctools.partial_972}} Resource localhost/_AnonymousVar3/N10tensorflow3VarE does not exist.
	 [[{{node random_rotation_1/stateful_uniform/RngReadAndSkip}}]]
	 [[random_rotation_1/stateful_uniform/RngReadAndSkip/_14]]
  (1) NOT_FOUND: {{function_node __inference_Dataset_map_classfunctools.partial_972}} Resource localhost/_AnonymousVar3/N10tensorflow3VarE does not exist.
	 [[{{node random_rotation_1/stateful_uniform/RngReadAndSkip}}]]
0 successful operations.
0 derived errors ignored. [Op:IteratorGetNext]
2021-12-13 16:10:16.949380: W tensorflow/core/kernels/data/generator_dataset_op.cc:107] Error occurred when finalizing GeneratorDataset iterator: FAILED_PRECONDITION: Python interpreter state is not initialized. The process may be terminated.
	 [[{{node PyFunc}}]]