Division by zero caused by mask operation

[Chenyang-1024]

If each pixel in the input image does not belong to the q-th class, then when generating the mask for masked attention, attn_mask[b, q, :] = True will be converted to attn_mask[b, q, :] = float('-inf') in nn.MultiheadAttention. Finally, when attn_mask is used for the Softmax(attn_mask, dim=-1) operation to calculate the attention map, the NaN caused by the divide by 0 error will appear. : (
This problem came up when I applied masked attention to my semantic segmentation task. : (

[q1556450920]

attn_mask[torch.where(attn_mask.sum(-1) == attn_mask.shape[-1])] = False Is there no such line?