some code cleaning

OpenNMT · May 21, 2024 · e9edb12 · e9edb12
1 parent 2c4eded
commit e9edb12
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/onmt/modules/multi_headed_attn.py b/onmt/modules/multi_headed_attn.py
@@ -487,7 +487,7 @@ def forward(
                         x = key_pad_mask
                         x = x.expand(-1, self.head_count // self.parallel_gpu, -1)
                         x = x.unsqueeze(3)
-                        x = x.expand(-1, -1, -1, 128)
+                        x = x.expand(-1, -1, -1, value.size(3))
                         value = value.masked_fill(x, 0)
 
                     self.layer_cache[1]["keys"] = key