Hi, Thank you for your reply! I am wondering the generate_square_masks can only generate square mask, but according to the nn.Transformer documentation, memory_mask requres shape as (T, S). So I am wondering do I need a customized function to generate it? Thank you!!