x = np.random.randn(150, 9)
y = np.random.randint(0, 10, 150)

x_train, y_train, x_val, y_val = split(x, y)

train_dataset = tf.data.Dataset.from_tensor_slices((x_train, tf.one_hot(y_train, depth=10)))
train_dataset = train_dataset.batch(32).repeat()

val_dataset = tf.data.Dataset.from_tensor_slices((x_val, tf.one_hot(y_val, depth=10)))
val_dataset = val_dataset.batch(32).repeat()