Can you recode so that you can input larger batches instead of batch size of 1?

This is my question: how to recode it to input the entire batch instead of calculating loss one by one.
If I try to input the batch to the criterion, I get:

ValueError: Expected target size (7, 35620), got torch.Size([7, 50])

I’m inputting torch.Size([7, 50, 35620]) as decoder outputs and torch.Size([7, 50]) as target indexes.
Where 50 is a batch_size, 35620 – vocabulary size (number of classes), and 7 is a number of predicted words.
It works when I pass it one by one (torch.Size([7, 35620]) with torch.Size([7])), and I don’t know how to adapt it for the batch input.
EDIT:
OMG, I tried “batch_first” way earlier and this didn’t work. But I didn’t try to put batch_size to the last position like torch.Size([7, 35620, 50]) and torch.Size([7, 50]). And this works now. Is this correct way to compute loss?
opt.zero_grad()
for i in range(iterations):
  ...
  loss = criterion_test(dec_outs.view(-1, vocab_size, batch_size), targets.view(-1, batch_size))
  loss.backward()
## after few iterations
opt.step()