The nn.Block should probably be renamed to nn.DecoderBlock as it's not very clear if it's supposed to be an encoder or a decoder block, also an option to disable the mask from attention. That said, very cool project.