Hi @albanD, I figured the nan source in the forward pass, It’s a masked softmax that uses -inf to mask the False values, but I guess I have many -infs that’s why it can return nan.
    unnorm.masked_fill_(emask,-float('inf'))
    attn = F.softmax(unnorm,dim=2)
    out = torch.bmm(attn,emb)

I tried the below line as alternative, but the values that should be masked have values near the others that shouldn’t. Any ideas here?
    unnorm.masked_fill_(emask,-1 * 1e-15)