TypeError: Couldn't cast array of type int64 while mapping the dataset

Hi!

It seems like the cast fails on the labels column - preprocess_data returns a single label per example, so the correct feature type for it would be Value("int64").