fix mask bias dtype in sdpa (#2407)

* fix mask bias dtype in sdpa * Update decoder.py
wenet-e2e · Mar 12, 2024 · 696d161 · 696d161
1 parent 983e86f
commit 696d161
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/wenet/transformer/decoder.py b/wenet/transformer/decoder.py
@@ -175,8 +175,8 @@ def forward(
         # tgt_mask: (B, L, L)
         tgt_mask = tgt_mask & m
         if self.use_sdpa:
-            tgt_mask = mask_to_bias(tgt_mask, tgt.dtype)
-            memory_mask = mask_to_bias(memory_mask, memory_mask.dtype)
+            tgt_mask = mask_to_bias(tgt_mask, memory.dtype)
+            memory_mask = mask_to_bias(memory_mask, memory.dtype)
 
         x, _ = self.embed(tgt)
         if self.gradient_checkpointing and self.training: