dataset_path = "/content/PetImages"
dataset = tf.keras.utils.image_dataset_from_directory(
    dataset_path,
    batch_size=32,
    image_size=(224, 224),  # Resize images
    shuffle=True
)

# Defining Albumentations Transform
transform = A.Compose([
    A.HorizontalFlip(p=0.5),
    A.RandomBrightnessContrast(p=0.2),
    A.Rotate(limit=30, p=0.5),
    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),  # Normalization
])

# Function to apply augmentation
def augment_image(image, label):
    img = image.numpy().astype(np.uint8)  # Convert Tensor to NumPy
    augmented = transform(image=img)["image"]
    return augmented.astype(np.float32), label  # Ensure float dtype

# Wrapping augmentation function for TF dataset
def tf_augment(image, label):
    img, lbl = tf.numpy_function(func=augment_image, inp=[image, label], Tout=[tf.float32, tf.int32])
    img.set_shape((224, 224, 3))
    return img, lbl

# Applying augmentation
augmented_dataset = dataset.map(tf_augment, num_parallel_calls=tf.data.AUTOTUNE)
augmented_dataset = augmented_dataset.batch(32).prefetch(tf.data.AUTOTUNE)

This is the input layer:

model.add(Conv2D(32, kernel_size=(3,3), padding='valid', activation='relu', input_shape=(224, 224, 3)))

And this is the error:

ValueError: Cannot take the length of shape with unknown rank.