Essentially, there is no restriction that says all batches must be the same size. But some of the code you are using may have been written with that assumption. Either you can track down the error and change the assumption (could take some time and thinking) or you could set drop_last in your DataLoader to True, which will make sure all batches are the same size.