Update modelling_hibial.py

modelling_hibial.py

@@ -180,7 +180,7 @@ class TransformerEncoderLayer(Module):
 
         # Extend mask
         # src_mask = src_mask.repeat(self.self_attn.num_heads, 1, 1)
-
+        
         # PreLayerNorm
         if self.norm_first:
 
@@ -334,6 +334,8 @@ class HiBiAlBert(Module):
         bialibi_utt_mask[enc_mask_utt.bool()] = float('-inf')
         bialibi_ct_mask[enc_mask_ct.bool()] = float('-inf')
 
+        src_key_padding_mask = torch.where(src_key_padding_mask, torch.tensor(float('-inf')), torch.tensor(0.0))
+        
         # memory = self.encoder(input_ids, mask=src_mask, src_key_padding_mask=src_key_padding_mask)
 
         # Encoding