I fixed the issue, thanks @ptrblck.
the mask was created using the following function:
(the rest of the code is quite big…)
    def make_mask(self, feature):
        return (torch.sum(
            torch.abs(feature),
            dim=-1
        ) == 0).unsqueeze(1).unsqueeze(2)

I now created the mask using not just the features, but from a zeros/ones ‘mask’ that I had created earlier when creating the sentence embedding matrix (and then putting that through this function).
Some of the embedding features should have been zero but were not probably because of in-between ops, that still messed up the softmax.
Hope that is enough info  .