So, that tensor shows what are the tokens predicted for each sentence.
In that example, it started with 101 which means start of sentence, and then it kept repeating 1996 which means “The”.
I’ve tried the same approach using Bart, but with no success.
Here’s an example:
from torch.utils.data import DataLoader

model = transformers.BartForConditionalGeneration.from_pretrained('facebook/bart-large').to(device)
tokenizer = transformers.BartTokenizer.from_pretrained('facebook/bart-large')

dl = DataLoader(dataset['train'], batch_size=8)

optimizer = torch.optim.AdamW(model.parameters())

epochs = 11
max_len = 50

data = next(iter(dl))

for e in range(epochs):

    en_tokenized_data = tokenizer.batch_encode_plus(
        data['translation']['en'],
        max_length=max_len, 
        truncation=True,
        pad_to_max_length=True, 
        return_attention_mask=True, 
        return_tensors='pt'
    )

    en_tokenized_data = en_tokenized_data.to(device)
  
    model.train()
    optimizer.zero_grad()

    result = model(
        input_ids=en_tokenized_data['input_ids'],
        attention_mask=en_tokenized_data['attention_mask'],
        decoder_input_ids=en_tokenized_data['input_ids'],
        decoder_attention_mask=en_tokenized_data['attention_mask'],
        labels=en_tokenized_data['input_ids']
    )

    loss, output = result[:2]
    print(f"Epoch: {e}, Loss: {loss.item()}")
    loss.backward()
    optimizer.step()

Before starting the training phase, it seems the model knows how to rebuild the original sentence:
Examples: 
 <s>ResResumption of the session</s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>
Original:
'Resumption of the session'

However, after 10 epochs, it starts generating garbage:
Examples: 
 <s> the the the the the the the the the the the the the the the the the the the

I think the issue lies in how to use the causal mask. I am not sure how to pass it to the model. It’s not an argument for the forward pass. Reading the docs I thought it was generated automatically by the model.
Or maybe I am trying to use the incorrect model for this task. I am pretty confused tbh.
Am I doing the training process correctly? In this I just wanted to learn how to reconstruct 8 sentences from my dataset.
Thanks!