Just to confirm, I can across a similar issue with the dataloader.
@smth @ptrblck, this seems to be an edge case where getting the last batch could result in some dimension mismatching, and it’s not ideal to have custom reshaping code to correct for this.  Is it possible to have a drop_the_last_batch option in the dataloader constructor?  That way all code downstream could assume that the batch size will be strictly consistent.
EDIT: I think I spoke too soon – looks like there is a drop_last option in the dataloader.  So that may help with the batch issues. https://pytorch.org/docs/stable/data.html 15