diff --git a/diffusion_planner/diffusion_planner/model/module/dit.py b/diffusion_planner/diffusion_planner/model/module/dit.py
index eb5b2f61f..65ccf956e 100644
--- a/diffusion_planner/diffusion_planner/model/module/dit.py
+++ b/diffusion_planner/diffusion_planner/model/module/dit.py
@@ -34,7 +34,7 @@ def __init__(self, dim=192, heads=6, dropout=0.1, mlp_ratio=4.0):
             in_features=dim, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0
         )
 
-    def forward(self, x, cross_c, y, attn_mask):
+    def forward(self, x, cross_c, y, attn_mask, cross_attn_mask):
         shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(
             y
         ).chunk(6, dim=2)
@@ -55,7 +55,16 @@ def forward(self, x, cross_c, y, attn_mask):
         modulated_x = modulate(self.norm2(x), shift_mlp, scale_mlp)
         x = x + gate_mlp * self.mlp1(modulated_x)
 
-        x = x + self.cross_attn(self.norm3(x), cross_c, cross_c, need_weights=False)[0]
+        x = (
+            x
+            + self.cross_attn(
+                self.norm3(x),
+                cross_c,
+                cross_c,
+                key_padding_mask=cross_attn_mask,
+                need_weights=False,
+            )[0]
+        )
         x = x + self.mlp2(self.norm4(x))
 
         return x
@@ -154,9 +163,10 @@ def forward(self, x, t, cross_c, neighbor_current_mask):
 
         ego_mask = torch.zeros((B, 1), dtype=torch.bool, device=x.device)
         attn_mask = torch.cat([ego_mask, neighbor_current_mask], dim=1)
+        cross_attn_mask = torch.all(cross_c == 0, dim=-1)
 
         for block in self.blocks:
-            x = block(x, cross_c, t, attn_mask)
+            x = block(x, cross_c, t, attn_mask, cross_attn_mask)
 
         x = self.final_layer(x, t)  # (B, P, output_dim)
         x = x.reshape(B, P, T, D)
diff --git a/diffusion_planner/diffusion_planner/model/module/encoder.py b/diffusion_planner/diffusion_planner/model/module/encoder.py
index ac63ce789..3ac4932a4 100644
--- a/diffusion_planner/diffusion_planner/model/module/encoder.py
+++ b/diffusion_planner/diffusion_planner/model/module/encoder.py
@@ -305,6 +305,7 @@ def forward(self, inputs):
         encoding_input = encoding_input + encoding_pos_result.view(B, self.token_num, -1)
 
         encoder_outputs = self.fusion(encoding_input, encoding_mask.view(B, self.token_num))
+        encoder_outputs = encoder_outputs.masked_fill(encoding_mask.view(B, self.token_num, 1), 0.0)
 
         return encoder_outputs