Lukas

Thanks for running the code and reporting this, and sorry it took me so long to reply. You're right: the printouts in the batched section came from an earlier run with different sentences, so they didn't match the code in the article. The code itself works; only the printed outputs were wrong. Your values are what the code produces, and they match the notebook.

I've updated the article. The outputs are now text instead of screenshots, taken from the notebook:

eos_idx:          tensor([ 7, 13, 19, 28, 37, 38])
eos_idx_expanded: tensor([ 0,  7, 13, 19, 28, 37, 38])
normalized_idx:   tensor([19,  7, 13, 19,  9, 18, 19])
reps:             tensor([7, 6, 6, 9, 9, 1])

Can you maybe share a full script to reproduce this error?

Yes, this sounds like the model just learns to not attend to earlier tokens when an eos token comes after. With a packed sequence mask you can enforce this explicitly by just masking out previous tokens.

Hey @shantanuagarwal , glad you enjoyed the article! Even though I havent tried it out myself you should be able to leverage pytorch flexattention api for this. Have a look at the tutorial here https://pytorch.org/blog/flexattention/. Section "Document Masking/Jagged Sequences" talks about these packed sequence masks.