input_ids is of shape (batch_size, max_len), i.e., torch.Size([16, 100]).
I’m not sure what you mean by context though.
This is the training routine including the output comparison:
  for d in data_loader:
     input_ids = d['enc_input_ids'].to(device)
     attention_mask = d['enc_attention_mask'].to(device)
     dec_input_ids = d['dec_input_ids'].to(device)
     dec_attention_mask = d['dec_attention_mask'].to(device)

     outputs = model(
        input_ids,
        attention_mask,
        dec_input_ids,
        dec_attention_mask,
        labels=dec_input_ids,
     )

    loss, logits = outputs.loss, outputs.logits

    decoded_output = tokenizer.batch_decode(torch.argmax(outputs.logits, dim=2).tolist(), skip_special_tokens=True)

    # Compare outputs...
    for i, output in enumerate(decoded_output)):
        if random.random() < 0.05: # only show a few random examples
            generated = model.generate(input_ids, decoder_start_token_id=model.config.pad_token_id)
            print("Forward pass output: ", output)
            print("Generate output: ", tokenizer.decode(generated[i], skip_special_tokens=True))

    loss.backward()
    nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    optimizer.step()
    scheduler.step()
    optimizer.zero_grad()