Batch size question

[QuintonQu]

Hi thanks to your good work. However, I found a problem. The code will raise error when the batch_size != 1. Is the batch size really useful in your code?
flow_np_x = flow_np[:, 0, :, :].reshape(self.height, self.width)