diff --git a/lighthouse/common/CIM.py b/lighthouse/common/CIM.py
index 71031bd..71530d7 100644
--- a/lighthouse/common/CIM.py
+++ b/lighthouse/common/CIM.py
@@ -700,7 +700,7 @@ def forward_post(self,
         # print(src_key_padding_mask[:, 76:].shape) # torch.Size([32, 26])
 
         qmask, kmask = src_key_padding_mask[:, 1:video_length + 1].unsqueeze(2), src_key_padding_mask[:, video_length + 1:].unsqueeze(1)
-        attn_mask = torch.matmul(qmask.float(), kmask.float()).bool().repeat(self.nhead, 1, 1)
+        attn_mask = torch.matmul(qmask.float(), kmask.float()).bool().repeat_interleave(self.nhead, dim=0)  # batch-major layout for nn.MultiheadAttention 3D attn_mask
         # print(attn_mask.shape)
         # print(attn_mask[0][0])
         # print(q.shape) 75 32 256
@@ -789,7 +789,7 @@ def forward_post(self,
         global_token, q, k, v = src[0].unsqueeze(0), pos_src[video_length + 1: ], pos_src[1:video_length + 1], src[1:video_length + 1]
 
         qmask, kmask = src_key_padding_mask[:, video_length + 1:].unsqueeze(2), src_key_padding_mask[:, 1:video_length + 1].unsqueeze(1)
-        attn_mask = torch.matmul(qmask.float(), kmask.float()).bool().repeat(self.nhead, 1, 1)
+        attn_mask = torch.matmul(qmask.float(), kmask.float()).bool().repeat_interleave(self.nhead, dim=0)  # batch-major layout for nn.MultiheadAttention 3D attn_mask
         # print(attn_mask.shape)
         # print(attn_mask[0][0])
         # print(q.shape) 23 32 256
diff --git a/lighthouse/common/cg_detr_transformer.py b/lighthouse/common/cg_detr_transformer.py
index 2a2340d..1a9bc2b 100644
--- a/lighthouse/common/cg_detr_transformer.py
+++ b/lighthouse/common/cg_detr_transformer.py
@@ -576,7 +576,7 @@ def forward_post(self,
 
         qmask, kmask = src_key_padding_mask[:, :video_length].unsqueeze(2), src_key_padding_mask[:,
                                                                                  video_length:].unsqueeze(1)
-        attn_mask = torch.matmul(qmask.float(), kmask.float()).bool().repeat(self.nhead, 1, 1)
+        attn_mask = torch.matmul(qmask.float(), kmask.float()).bool().repeat_interleave(self.nhead, dim=0)  # batch-major layout for nn.MultiheadAttention 3D attn_mask
 
         # - key_padding_mask: :math:`(S)` or :math:`(N, S)` where N is the batch size, S is the source sequence length.
         #   If a FloatTensor is provided, it will be directly added to the value.
diff --git a/lighthouse/common/qd_detr_transformer.py b/lighthouse/common/qd_detr_transformer.py
index 363658c..3697ad1 100755
--- a/lighthouse/common/qd_detr_transformer.py
+++ b/lighthouse/common/qd_detr_transformer.py
@@ -476,7 +476,7 @@ def forward_post(self,
         # print(src_key_padding_mask[:, 76:].shape) # torch.Size([32, 26])
 
         qmask, kmask = src_key_padding_mask[:, 1:video_length + 1].unsqueeze(2), src_key_padding_mask[:, video_length + 1:].unsqueeze(1)
-        attn_mask = torch.matmul(qmask.float(), kmask.float()).bool().repeat(self.nhead, 1, 1)
+        attn_mask = torch.matmul(qmask.float(), kmask.float()).bool().repeat_interleave(self.nhead, dim=0)  # batch-major layout for nn.MultiheadAttention 3D attn_mask
         # print(attn_mask.shape)
         # print(attn_mask[0][0])
         # print(q.shape) 75 32 256
diff --git a/lighthouse/common/taskweave_transformer.py b/lighthouse/common/taskweave_transformer.py
index cd7363b..87369b2 100644
--- a/lighthouse/common/taskweave_transformer.py
+++ b/lighthouse/common/taskweave_transformer.py
@@ -591,7 +591,7 @@ def forward_post(self,
         q, k, v = pos_src[:video_length], pos_src[video_length:], src[video_length:]
 
         qmask, kmask = src_key_padding_mask[:, :video_length].unsqueeze(2), src_key_padding_mask[:, video_length :].unsqueeze(1)
-        attn_mask = torch.matmul(qmask.float(), kmask.float()).bool().repeat(self.nhead, 1, 1)
+        attn_mask = torch.matmul(qmask.float(), kmask.float()).bool().repeat_interleave(self.nhead, dim=0)  # batch-major layout for nn.MultiheadAttention 3D attn_mask
 
         src2 = self.self_attn(q, k, value=v, attn_mask=attn_mask,
                               key_padding_mask=src_key_padding_mask[:, video_length:])[0]
diff --git a/lighthouse/common/tr_detr_transformer.py b/lighthouse/common/tr_detr_transformer.py
index d4141d2..d3d8b58 100644
--- a/lighthouse/common/tr_detr_transformer.py
+++ b/lighthouse/common/tr_detr_transformer.py
@@ -434,7 +434,7 @@ def forward_post(self,
         # print(src_key_padding_mask[:, 76:].shape) # torch.Size([32, 26])
         # print(src_key_padding_mask[:, 1:video_length+1].shape) # torch.Size([32, 97])  切片会保留维度torch.Size([32, 75])，只有坐标便不会
         qmask, kmask = src_key_padding_mask[:, :video_length].unsqueeze(2), src_key_padding_mask[:, video_length:].unsqueeze(1)
-        attn_mask = torch.matmul(qmask.float(), kmask.float()).bool().repeat(self.nhead, 1, 1)
+        attn_mask = torch.matmul(qmask.float(), kmask.float()).bool().repeat_interleave(self.nhead, dim=0)  # batch-major layout for nn.MultiheadAttention 3D attn_mask
         # print(attn_mask.shape)
         # print(attn_mask[0][0])
         # print(q.shape) 75 32 256
@@ -525,7 +525,7 @@ def forward_post(self,
         # print(src_key_padding_mask[:, 76:].shape) # torch.Size([32, 26])
 
         qmask, kmask = src_key_padding_mask[:, 1:video_length + 1].unsqueeze(2), src_key_padding_mask[:, video_length + 1:].unsqueeze(1)
-        attn_mask = torch.matmul(qmask.float(), kmask.float()).bool().repeat(self.nhead, 1, 1)
+        attn_mask = torch.matmul(qmask.float(), kmask.float()).bool().repeat_interleave(self.nhead, dim=0)  # batch-major layout for nn.MultiheadAttention 3D attn_mask
         # print(attn_mask.shape)
         # print(attn_mask[0][0])
         # print(q.shape) 75 32 256