When I set the batch size to 1, the code works,  but when I set the batch size to greater than 1, it will throw an error