Upload model_Custm.py

model_Custm.py

@@ -288,6 +288,7 @@ class Wildnerve_tlm01(nn.Module, AbstractModel):
         labels=None,
         src=None,
         tgt=None,
+        src_mask: Optional[torch.Tensor] = None,          # added
         src_key_padding_mask=None,
         tgt_key_padding_mask=None,
         memory_key_padding_mask=None,
@@ -319,7 +320,7 @@ class Wildnerve_tlm01(nn.Module, AbstractModel):
                 key = query
                 value = query
                 
-                # IMPORTANT: Initialize src_mask if it's None
+                # CRITICAL: Initialize src_mask if it's None
                 if src_mask is None and src is not None:
                     # Create a default mask that allows all tokens to attend to all other tokens
                     src_seq_len = src.size(1)