Batch inter_chain_mask shape mismatch in multi-head attention #6

Hello! When using batched inputs, inter_chain_mask has shape [B, L, L], while attention scores are reshaped to [B * H, L, L].
In pplm/multihead_attention.py line 362-366, the current code repeats the mask like this:

attn_intra[inter_chain_mask.repeat(attn_intra.shape[0], 1, 1).bool()] = 0
attn_inter[(1 - inter_chain_mask.repeat(attn_intra.shape[0], 1, 1)).bool()] = 0

attn_weights = attn_intra + self.inter_attn_weight[:, None, None] * attn_inter
attn_weights = attn_weights + inter_chain_mask[None]

This works for single-sample input where inter_chain_mask is [L, L], but fails for batched input.