The final thing you want to do at the end of your pipeline is to produce a batch of data ready to send to your GPU for training. This is done simply with the batch method and passing in the size of batch you want:

train_dataset = train_dataset.batch(128)  # Take a batch of 128 from the dataset. 

The size of the batch is an important parameter when trying to make our pipeline as efficient as possible; having it as large as possible may not always be best. For example, if you have lots of preprocessing steps on your images, then your GPU may be standing idle while your CPU does preprocessing for a large batch of images, getting them ready.