Mismatch in batch size

Also posted here with follow-up.