The logits have (8, 2, 32128) shape, and the labels (8,2). After applying the view and size functions the final shapes are (16, 32128) for the logits and (16, 2) for the labels:
logits.view(-1, logits.size(-1)), lm_labels.view(-1)

Any toughts?