Chapter 3 MultiHeadAttention boolean mask

[rayed-therap]

# Original mask truncated to the number of tokens and converted to boolean
mask_bool = self.mask.bool()[:num_tokens, :num_tokens]     
# Unsqueeze the mask twice to match dimensions        
mask_unsqueezed = mask_bool.unsqueeze(0).unsqueeze(0)
# Use the unsqueezed mask to fill attention scores
attn_scores.masked_fill_(mask_unsqueezed, -torch.inf)

Are the double unsqueezes necessary? Won't mask_bool broadcast in attn_scores.masked_fill_ as the last two dimensions are equal to num_tokens?

[rasbt]

You are right, the double-unsqueeze is redundant. This must have been a leftover from some previous code I had where this was necessary. Thanks!!