I also tried add that tensor to model’s input. But with this case, when I use 2 gpus, the result is just about half of one batch, (64 x ...) tensor returned, just half of one batch (128).
0/5120 [00:00<? ?it/s] Traceback (most recent call last):
  File "main.py", line 234, in <module>
    trainer.train()
  File "main.py", line 116, in train
    loss_lm = self.criterion(logits_lm.transpose(1, 2), batch.masked_tokens.transpose(0,1)) # for masked LM
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/loss.py", line 932, in forward
    ignore_index=self.ignore_index, reduction=self.reduction)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/functional.py", line 2317, in cross_entropy
    return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/functional.py", line 2113, in nll_loss
    .format(input.size(0), target.size(0)))
ValueError: Expected input batch_size (64) to match target batch_size (128).

and below is training loop.
        pos = torch.arange(513, dtype=torch.long, requires_grad=False).unsqueeze(0).to(device=torch.device('cuda'))
        for epoch in range(max_epoch):
            loss_sum, acc_sum, len_batch_sum = 0., 0., 0.
            ds_iter.init_epoch()
            tr_total = math.ceil(total_len / self.batch_size)
            tq_iter = tqdm(enumerate(ds_iter), total=tr_total, miniters=min_iters, unit_scale=self.batch_size,
                           bar_format='{n_fmt}/{total_fmt} [{elapsed}<{remaining} {rate_fmt}] {desc}')

            self.model.train()
            print('epoch starts')
            for i, batch in tq_iter:
                self.model.zero_grad()
                print('batch starts')
                device = torch.device('cuda')
                print(device)
                print(f'batch.input_ids device : {batch.input_ids.device}, batch.segment_ids : {batch.segment_ids.device}, batch.masekd_pos : {batch.masked_pos.device}')
                print(f'batch.input_ids shape : {batch.input_ids.shape}, batch.segment_ids : {batch.segment_ids.shape}, batch.masekd_pos : {batch.masked_pos.shape}, pos : {pos.shape}')
                logits_lm, logits_clsf = self.model(batch.input_ids.transpose(0,1).to(device=device), batch.segment_ids.transpose(0,1).to(device=device), batch.masked_pos.transpose(0,1).to(device=device), pos.to(device=device))
                print(f'logits_lm, logits_clsf shape : {logits_lm.shape}, {logits_clsf.shape}')

and these are logs.
epoch starts
batch starts
cuda
batch.input_ids device : cuda:0, batch.segment_ids : cuda:0, batch.masekd_pos : cuda:0
batch.input_ids shape : torch.Size([513, 128]), batch.segment_ids : torch.Size([513, 128]), batch.masekd_pos : torch.Size([5, 128]), pos : torch.Size([1, 513])
pos device: cuda:0 # log from Embedding layer in the model
pos device: cuda:0, x device : cuda:0, seg device: cuda:0 # log from Embedding layer in the model
logits_lm, logits_clsf shape : torch.Size([64, 5, 6015]), torch.Size([64, 2])

logits_lm device: cuda:0, batch target device: cuda:0
0/5120 [00:00<? ?it/s] Traceback (most recent call last):
  File "main.py", line 236, in <module>
    trainer.train()
  File "main.py", line 118, in train
    loss_lm = self.criterion(logits_lm.transpose(1, 2), batch.masked_tokens.transpose(0,1)) # for masked LM
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/loss.py", line 932, in forward
    ignore_index=self.ignore_index, reduction=self.reduction)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/functional.py", line 2317, in cross_entropy
    return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/functional.py", line 2113, in nll_loss
    .format(input.size(0), target.size(0)))
ValueError: Expected input batch_size (64) to match target batch_size (128).