Switch dataset code from keras sequence to tf.data

[soi]

Looked into it, seems to be very complicated to get right with our data and setup. Probably not worth the effort.

[alisandra]

Ah well, thanks for checking