Using BartForConditionalGeneration with a batch_size = 2 … all the inputs look right but when I examine the logits the shape is torch.Size([2, 1, 50264]) 
The inputs are:
x['input_ids'].shape => torch.Size([2, 256])
x['attention_mask'].shape => torch.Size([2, 256])
x['decoder_input_ids'].shape => torch.Size([2, 68])

What may I be doing wrong? Or is there a bug in the model?