Yes, I also faced this problem.
Traceback (most recent call last):                                                                 | 0/765 [00:00<?, ?it/s]
outputs = model(**inputs)

File “/usr/local/lib/python3.6/site-packages/torch/nn/modules/module.py”, line 532, in call
result = self.forward(*input, **kwargs)
File “/usr/local/lib/python3.6/site-packages/torch/nn/parallel/data_parallel.py”, line 152, in forward
outputs = self.parallel_apply(replicas, inputs, kwargs)
File “/usr/local/lib/python3.6/site-packages/torch/nn/parallel/data_parallel.py”, line 162, in parallel_apply
return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
File “/usr/local/lib/python3.6/site-packages/torch/nn/parallel/parallel_apply.py”, line 85, in parallel_apply
output.reraise()
File “/usr/local/lib/python3.6/site-packages/torch/_utils.py”, line 394, in reraise
raise self.exc_type(msg)
ValueError: Caught ValueError in replica 0 on device 0.
Original Traceback (most recent call last):
File “/usr/local/lib/python3.6/site-packages/torch/nn/parallel/parallel_apply.py”, line 60, in _worker
output = module(*input, **kwargs)
File “/usr/local/lib/python3.6/site-packages/torch/nn/modules/module.py”, line 532, in call
result = self.forward(*input, **kwargs)
File “/workspace/git/transformers/src/transformers/modeling_bert.py”, line 1522, in forward
loss = loss_fct(logits, labels)
File “/usr/local/lib/python3.6/site-packages/torch/nn/modules/module.py”, line 532, in call
result = self.forward(*input, **kwargs)
File “/usr/local/lib/python3.6/site-packages/torch/nn/modules/loss.py”, line 916, in forward
ignore_index=self.ignore_index, reduction=self.reduction)
File “/usr/local/lib/python3.6/site-packages/apex/amp/wrap.py”, line 28, in wrapper
return orig_fn(*new_args, **kwargs)
File “/usr/local/lib/python3.6/site-packages/torch/nn/functional.py”, line 2021, in cross_entropy
return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)
File “/usr/local/lib/python3.6/site-packages/apex/amp/wrap.py”, line 28, in wrapper
return orig_fn(*new_args, **kwargs)
File “/usr/local/lib/python3.6/site-packages/torch/nn/functional.py”, line 1836, in nll_loss
.format(input.size(0), target.size(0)))
ValueError: Expected input batch_size (2) to match target batch_size (4).

After modified the batch data, but forgot to multiple * max(1, args.n_gpu) on the sample number dimension.
@ptrblck   So many thanks!