I was thinking to remove this whole if statement everywhere. Since in debug mode the data batch size is set to 4 or 1, but not the actual model batch size. What do you think @goodfeli ?