Problem solved thanks to @Deepti_Prasad
If you do get the same error, make sure you re-check your train_val_dataset function, and that the resulting dataset has the desired shape, beyond cardinality.
In my case, the unittests did not flag my mistake.