From 4f7d5f0d209f929f00047beaedc73ae09cf3371b Mon Sep 17 00:00:00 2001
From: seasideGargantua <102777023+seasideGargantua@users.noreply.github.com>
Date: Wed, 20 May 2026 10:51:10 +0800
Subject: [PATCH] Fix prompt padding mask per sample

---
 .../model/dreamzero/action_head/wan_flow_matching_action_tf.py  | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/groot/vla/model/dreamzero/action_head/wan_flow_matching_action_tf.py b/groot/vla/model/dreamzero/action_head/wan_flow_matching_action_tf.py
index b7cd806f..e20abc26 100644
--- a/groot/vla/model/dreamzero/action_head/wan_flow_matching_action_tf.py
+++ b/groot/vla/model/dreamzero/action_head/wan_flow_matching_action_tf.py
@@ -544,7 +544,7 @@ def encode_prompt(self, input_ids, attention_mask):
         prompt_emb = self.text_encoder(input_ids, attention_mask)
         prompt_emb = prompt_emb.clone().to(dtype=torch.bfloat16)
         for i, v in enumerate(seq_lens):
-            prompt_emb[:, v:] = 0
+            prompt_emb[i, v:] = 0
         return prompt_emb
 
     def _ensure_vae_on_device(self, ref_tensor):