Option to pass in memory_mask

Hi Phil,

Can we have the option to pass in memory_mask, just like in Official Pytorch Transformers?

Thanks.

[lucidrains]

ok, i’ve deployed it in 0.17.10, let me know if that works for you!

import torch
from x_transformers import XTransformer

model = XTransformer(
    dim = 512,
    enc_num_tokens = 256,
    enc_depth = 1,
    enc_heads = 8,
    enc_max_seq_len = 1024,
    dec_num_tokens = 256,
    dec_depth = 1,
    dec_heads = 8,
    dec_max_seq_len = 1024
)

src = torch.randint(0, 256, (1, 1024))
src_mask = torch.ones_like(src).bool()
tgt = torch.randint(0, 256, (1, 1024))
tgt_mask = torch.ones_like(tgt).bool()

attn_mask = ~torch.eye(1024).bool()

loss = model(src, tgt, src_mask = src_mask, tgt_mask = tgt_mask, src_attn_mask = attn_mask) # (1, 1024, 512)
loss.backward()

[tumbleintoyourheart]

I did some quick tests and it gave instant better results with proper masking thanks to the new feature.

Thanks for your quick responses, it’s very nice of you to even write down some sample codes for me. Really appreciated.