for batch in range(num_data // batch_size + (num_data % batch_size > 0)):
...First get the number of full batches (num_data // batch_size), then add 1 if there are remaining data points (num_data % batch_size > 0).
Number of Batches
Correct iteration count including the remainder batch.
2020