How can I split output from the tf.data.Dataset?

Question

I have some .npy data whose shapes are [150~180, 480~512, 480~512], extracted from MRI image.

So I used some functions to refine the dataset and convert them as tf.da