For some reason this only works if the batch size evenly divides into the data. This seems like a Keras issue, any ideas on how to fix?