The issue was that an initialization was needed for the RandomFlip, RandomRotation layers.
This is the code that solves the issue (credit to M.Innat from StackOverflow):
def data_gen():
    for i in range(10):
        x = np.random.random(size=(80, 80, 3)) * 255  # rgb image
        x = x.astype('uint8')
        y = np.random.random(size=(40, 40, 1)) * 255  # downsized mono image
        y = y.astype('uint8')
        yield x, y


class Augment(tf.keras.layers.Layer):
    def __init__(self, seed=42):
        super().__init__()
        self.flip_a = tf.keras.layers.RandomFlip(mode="horizontal", seed=seed)
        self.flip_b = tf.keras.layers.RandomFlip(mode="horizontal", seed=seed)

        self.rot_a = tf.keras.layers.RandomRotation(factor=1.0,
                                 fill_mode='constant', seed=seed)
        self.rot_b = tf.keras.layers.RandomRotation(factor=1.0, 
                                fill_mode='constant', seed=seed)
        
    def call(self, inputs, labels):
        x = self.flip_a(inputs)
        x = self.rot_a(x)

        y = self.flip_b(labels)
        y = self.rot_b(y)
        return x, y


def preprocess(image, label, cropped_image_size, cropped_label_size):
    x = image
    y = label
    x_size = cropped_image_size
    y_size = cropped_label_size

    x = tf.cast(x, dtype=tf.float32)
    y = tf.cast(y, dtype=tf.float32)

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

    x = tf.cast(x, dtype=tf.uint8)
    y = tf.cast(y, dtype=tf.uint8)
    return x, y


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

crop_only_fn = functools.partial(preprocess, 
                                 cropped_image_size=50,  
                                 cropped_label_size=25)

# This works
crop_dataset = dataset.map(crop_only_fn)
x, y = next(iter(crop_dataset))
x.shape, y.shape
(TensorShape([50, 50, 3]), TensorShape([25, 25, 1]))

train_preprocess_fn = functools.partial(preprocess, 
                                        cropped_image_size=50,
                                        cropped_label_size=25)
train_dataset = dataset.map(train_preprocess_fn)
train_dataset = train_dataset.map(Augment()) # < calling now.
x, y = next(iter(train_dataset))
x.shape, y.shape
(TensorShape([50, 50, 3]), TensorShape([25, 25, 1]))