Batch_size does NOT equal after dataparallel forward

Hi, I trained a model on multi gpu using dataparallel, but at last batch, it got ValueError: Expected input batch_size to match target batch_size, anyone knows why?

model = torch.nn.DataParallel(model, device_ids=[0,1,2])
criterion = torch.nn.CrossEntropyLoss()
inputs = Batch.inputs.cuda()
targets = Batch.targets.cuda()
assert inputs.size(0) == targets.size(0)
logits = model(inputs)
loss = criterion(logits, targets)

Could you print the shapes of the tensors yielding this error?
To get rid of this error, you could specify drop_last=True in your DataLoader for now.

I guess the last batch might have a single sample which might get squeezed somewhere, so that a size mismatch occurs.

the shape of TARGETS tensor in last yield batch is (12), and shape of INPUTS tensor is (12,128,512). But the shape of LOGITS tensor is (8, 311)