Fix mask dimensions: unsqueeze 3D->4D for OpenNMT-py 3.x attention

molscribe/transformer/decoder.py

@@ -172,6 +172,13 @@ class TransformerDecoderLayerBase(nn.Module):
             layer_cache["memory_keys"] = cache_dict.get("keys")
             layer_cache["memory_values"] = cache_dict.get("values")
 
+    @staticmethod
+    def _expand_mask(mask):
+        """Expand 3D mask (B, 1, L) to 4D (B, 1, 1, L) for OpenNMT-py 3.x."""
+        if mask is not None and mask.dim() == 3:
+            return mask.unsqueeze(1)
+        return mask
+
     def _forward_self_attn(self, inputs_norm, dec_mask, layer_cache, step):
         if isinstance(self.self_attn, MultiHeadedAttention):
             # OpenNMT-py 3.x: layer_cache and attn_type are instance attributes
@@ -181,7 +188,7 @@ class TransformerDecoderLayerBase(nn.Module):
                 inputs_norm,
                 inputs_norm,
                 inputs_norm,
-                mask=dec_mask,
+                mask=self._expand_mask(dec_mask),
             )
             self._from_onmt3_cache(self.self_attn.layer_cache, layer_cache, "self")
             return result
@@ -309,7 +316,7 @@ class TransformerDecoderLayer(TransformerDecoderLayerBase):
             memory_bank,
             memory_bank,
             query_norm,
-            mask=src_pad_mask,
+            mask=self._expand_mask(src_pad_mask),
         )
         self._from_onmt3_cache(self.context_attn.layer_cache, layer_cache, "context")
         output = self.feed_forward(self.drop(mid) + query)