i am trying to train mobilevit using this bee dataset from tensorflow and i keep getting this error
TypeError: Expected any non-tensor type, but got a tensor instead.

this is the code i wrote, everything works except this piece of code, is the problem with the dataset ?

train_dataset, val_dataset = tfds.load(
    "bee_dataset/bee_dataset_300", split=["train[:90%]", "train[90%:]"], as_supervised=True
)

num_train = train_dataset.cardinality()
num_val = val_dataset.cardinality()
print(f"Number of training examples: {num_train}")
print(f"Number of validation examples: {num_val}")

train_dataset = prepare_dataset(train_dataset, is_training=True)
val_dataset = prepare_dataset(val_dataset, is_training=False)

i tried another dataset and it worked just fine but the error persists for this one