Using albumentation's augmentation in tensorflow dataset API is giving this error : Incompatible shapes expected [?,224,224,3] but got [8,1,224,224,3]

Question

I am getting this error while trying to augment the images using the albumentations library which uses tf.numpy_function to wrap the python function for augmentation in tens