Added MHA and related nodes.

modules/pytorch_wrapper/pt_train_classification_transformer_model.py

@@ -24,6 +24,9 @@ class PtTrainClassificationTransformerModel:
             classification_metrics (bool): If True and if val_loader is set, print out classification metrics.  This is only valid for classification model training.  
             val_loader (PTDATALOADER): DataLoader for validation dataset (optional).  
             scheduler (PTLRSCHEDULER): Learning rate scheduler.  
+            unpack_input (Optional[bool]): If True and the input is a list, unpack it when calling the model.  
+            For example, if the dataloader returns [a, b] and your model's forward is defined as forward(a, b), set this to True.
+
 
     category: PyTorch wrapper - Training
     """
@@ -51,7 +54,8 @@ class PtTrainClassificationTransformerModel:
             },
             "optional": {
                 "scheduler": ("PTLRSCHEDULER", {}),
-                "val_loader": ("PTDATALOADER", {})
+                "val_loader": ("PTDATALOADER", {}),
+                "unpack_input":  ("BOOLEAN", {"default":False})
             }
         }
 
@@ -76,7 +80,8 @@ class PtTrainClassificationTransformerModel:
           output_best_val_model:bool,
           classification_metrics:bool,
           scheduler: Optional[Union[torch.optim.lr_scheduler._LRScheduler, torch.optim.lr_scheduler.ReduceLROnPlateau]] = None,
-          val_loader: Optional[torch.utils.data.DataLoader] = None) -> Tuple[nn.Module, torch.Tensor, torch.Tensor]:
+          val_loader: Optional[torch.utils.data.DataLoader] = None,
+          unpack_input: Optional[bool] = False) -> Tuple[nn.Module, torch.Tensor, torch.Tensor]:
         """
         Trains an Transformer model.
 
@@ -93,7 +98,7 @@ class PtTrainClassificationTransformerModel:
             classification_metrics (bool): If True and if val_loader is set, print out classification metrics.  This is only valid for classification model training.
             val_loader (Optional[torch.utils.data.DataLoader]): DataLoader for validation dataset (optional).
             scheduler (Optional[Union[torch.optim.lr_scheduler._LRScheduler, torch.optim.lr_scheduler.ReduceLROnPlateau]]): Learning rate scheduler.
-
+            unpack_input (Optional[bool]): Unpack input (x) if input is a list.
         Returns:
             Tuple[torch.nn.Module, torch.Tensor, torch.Tensor]: A tuple containing the trained model, 
             training loss history, and validation loss history.
@@ -135,7 +140,10 @@ class PtTrainClassificationTransformerModel:
                         y = y.to("cuda")
                     optimizer.zero_grad()
 
-                    y_hat = model(x)
+                    if unpack_input:
+                        y_hat = model(*x)
+                    else:
+                        y_hat = model(x)
                     if y.dim() == 1:
                         y = torch.unsqueeze(y, dim=-1)
                         y = y.to(torch.float32)
@@ -167,7 +175,10 @@ class PtTrainClassificationTransformerModel:
                                     x = x.to("cuda")
                                 y = y.to("cuda")
 
-                            y_hat = model(x)
+                            if unpack_input:
+                                y_hat = model(*x)
+                            else:
+                                y_hat = model(x)
                             if y.dim() == 1:
                                 y = torch.unsqueeze(y, dim=-1)
                                 y = y.to(torch.float32)

modules/pytorch_wrapper/ptn_chained_model.py

@@ -44,7 +44,7 @@ class PtnChainedModelDef(nn.Module):
             torch.Tensor: Output tensor after passing through `model_a`, `model_b`, and optionally `closure`.
         """
         x = self.model_a(inputs)
-        x = self.model_b(x)  # Fixed the incorrect use of `inputs` instead of `x`
+        x = self.model_b(x)
         if self.closure:
             x = self.closure(x)
         return x

modules/pytorch_wrapper/ptn_chained_model_with_attention_mask.py

@@ -0,0 +1,128 @@
+from typing import Any, Dict, Tuple, Callable, Optional
+import torch
+import torch.nn as nn
+
+
+class PtnChainedModelWithAttentionMaskDef(nn.Module):
+    """
+    A chained model that sequentially applies `model_a` and `model_b` optionally followed by a closure function. The model also takes attention mask for input.
+
+    ### Fields:  
+    - `model_a`: First model in the chain.
+    - `model_b`: Second model in the chain.
+    - `model_a_mask_req`: True if model_a requires attention_mask for input. For example, Set True if you plug in the `Ptn Multihead Attention` node.
+    - `model_b_mask_req`: True if model_b requires attention_mask for input
+    - `closure`: Optional differentiable function to be called at the end of forward.
+
+    pragma: skip_doc
+    """
+
+    def __init__(self,
+                 model_a: nn.Module,
+                 model_b: nn.Module,
+                 model_a_mask_req: bool,
+                 model_b_mask_req: bool,
+                 closure: Optional[Callable] = None) -> None:
+        super().__init__()
+        self.model_a = model_a
+        self.model_b = model_b
+        self.model_a_mask_req = model_a_mask_req
+        self.model_b_mask_req = model_b_mask_req
+        self.closure = closure
+
+    def forward(self, inputs: torch.Tensor, mask:torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass through the chained models.
+
+        Args:
+            inputs (torch.Tensor): Input tensor.
+            mask (torch.Tensor): Mask tensor.
+
+        Returns:
+            torch.Tensor: Output tensor after passing through `model_a`, `model_b`, and optionally `closure`.
+        """
+        if self.model_a_mask_req:
+            x = self.model_a(inputs, mask)
+        else:
+            x = self.model_a(inputs)
+
+        if self.model_b_mask_req:
+            x = self.model_b(x, mask)
+        else:
+            x = self.model_b(x)
+        if self.closure:
+            x = self.closure(x)
+        return x
+
+
+class PtnChainedModelWithAttentionMask:
+    """
+    A chained model that sequentially applies `model_a` and `model_b`,
+    optionally followed by a closure function.
+
+    ### Fields:  
+    - `model_a`: First model in the chain.
+    - `model_b`: Second model in the chain.
+    - `model_a_mask_req`: True if model_a requires attention_mask for input
+    - `model_b_mask_req`: True if model_b requires attention_mask for input
+    - closure: Optional differentiable function to be called at the end of forward.
+
+    category: PyTorch wrapper - Model
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the expected input types.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required and optional input types.
+        """
+        return {
+            "required": {
+                "model_a": ("PTMODEL", {}),
+                "model_b": ("PTMODEL", {}),
+                "model_a_mask_req": ("BOOLEAN", {"default": False}),
+                "model_b_mask_req": ("BOOLEAN", {"default": False}),
+            },
+            "optional": {
+                "closure": ("PTCALLABLE", {})
+            }
+        }
+
+    RETURN_TYPES: Tuple[str, ...] = ("PTMODEL",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Training"
+
+    @classmethod
+    def IS_CHANGED(cls, **kw):
+        return float("NaN")
+
+    def f(self,
+          model_a: nn.Module,
+          model_b: nn.Module,
+          model_a_mask_req: bool,
+          model_b_mask_req: bool,
+          closure: Optional[Callable] = None) -> Tuple[nn.Module]:
+        """
+        Constructs a chained model with optional closure.
+
+        Args:
+            model_a (nn.Module): First model in the chain.
+            model_b (nn.Module): Second model in the chain.
+            model_a_mask_req (bool): Whether model_a needs attention mask.
+            model_b_mask_req (bool): Whether model_b needs attention mask.
+            closure (Callable, optional): A function applied after `model_b`.
+
+        Returns:
+            Tuple[nn.Module]: A tuple containing the instantiated model.
+        """
+        with torch.inference_mode(False):
+            model = PtnChainedModelWithAttentionMaskDef(
+                model_a=model_a,
+                model_b=model_b,
+                model_a_mask_req=model_a_mask_req,
+                model_b_mask_req=model_b_mask_req,
+                closure=closure
+            )
+        return (model,)

modules/pytorch_wrapper/ptn_embedding.py

@@ -26,7 +26,7 @@ class PtnEmbedding:
         return {
             "required": {
                 "num_embeddings": ("INT", {"default": 10000, "min": 1, "max": 1e6}),
-                "embedding_size": ("INT", {"default": 100, "min": 1, "max": 1e6}),
+                "embedding_dim": ("INT", {"default": 100, "min": 1, "max": 1e6}),
                 "padding_idx": ("INT", {"default": -1, "min": -1, "max": 1e6}),
             }
         }

modules/pytorch_wrapper/ptn_masked_mean_pooling.py

@@ -0,0 +1,77 @@
+from typing import Any, Dict, Tuple
+import torch
+import torch.nn as nn
+
+class PtnMaskedMeanPoolingDef(nn.Module):
+    """
+    Constructs a masked mean pooling layer.
+
+    pragma: skip_doc
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, inputs: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        """
+        Computes the mean of the input excluding the padding values using the given mask.
+
+        Args:
+            inputs (Tensor): 3D input tensor in (Batch, Sequence, Features) format.
+            mask (Tensor): 2D input tensor in (Batch, Sequence) format. Padding is indicated by 0; valid tokens are indicated by 1.
+
+        Returns:
+            torch.Tensor: Output tensor.
+        """
+        if mask.dim() != 2:
+            raise ValueError("Mask is not rank 2")
+        if inputs.dim() != 3:
+            raise ValueError("Inputs is not rank 2")
+        
+        maskf = mask.float()  # (B, Seq)
+        x = (inputs * maskf.unsqueeze(-1)).sum(dim=1) / maskf.sum(dim=1, keepdim=True).clamp(min=1e-9)  # (B, Feature)
+        if x.dim() != 2:
+            raise ValueError("Outputs is not rank 2")
+        return x
+
+
+class PtnMaskedMeanPooling:
+    """
+    Constructs a masked mean pooling layer.
+    Computes the mean of the input excluding the padding values using the given mask.
+
+    category: PyTorch wrapper - Model
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types for the function.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required input types.
+        """
+        return {
+        }
+
+    RETURN_TYPES: Tuple[str] = ("PTMODEL",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Training"
+
+    @classmethod
+    def IS_CHANGED(cls, **kw):
+        return float("NaN")
+
+    def f(self) -> Tuple[nn.Module]:
+        """
+        Constructs a masked mean pooling layer.
+        Computes the mean of the input excluding the padding values using the given mask.
+
+        Returns:
+            Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
+        """
+ 
+        with torch.inference_mode(False):
+            model = PtnMaskedMeanPoolingDef()
+
+        return (model,)

modules/pytorch_wrapper/ptn_model_with_closure.py

@@ -0,0 +1,87 @@
+from typing import Any, Dict, Tuple, Callable, Optional
+import torch
+import torch.nn as nn
+
+
+class PtnModelWithClosureDef(nn.Module):
+    """
+    A model that is followed by a closure function.
+
+    ### Fields:  
+    - `model`: The model in the chain.
+    - `closure`: Differentiable function to be called at the end of forward.
+
+    pragma: skip_doc
+    """
+
+    def __init__(self,
+                 model: nn.Module,
+                 closure: Callable) -> None:
+        super().__init__()
+        self.model = model
+        self.closure = closure
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        """
+        A model that is followed by a closure function.
+
+        Args:  
+        - `model`: The model in the chain.
+        - `closure`: Differentiable function to be called at the end of forward.
+
+        Returns:
+            torch.Tensor: Output tensor after passing through `model` and `closure`.
+        """
+        x = self.model(inputs)
+        x = self.closure(x)
+        return x
+
+
+class PtnModelWithClosure:
+    """
+    A model that is followed by a closure function.
+
+    ### Fields:  
+    - `model`: The model in the chain.
+    - `closure`: Differentiable function to be called at the end of forward.    category: PyTorch wrapper - Model
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the expected input types.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required and optional input types.
+        """
+        return {
+            "required": {
+                "model": ("PTMODEL", {}),
+                "closure": ("PTCALLABLE", {})
+            }
+        }
+
+    RETURN_TYPES: Tuple[str, ...] = ("PTMODEL",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Training"
+
+    @classmethod
+    def IS_CHANGED(cls, **kw):
+        return float("NaN")
+
+    def f(self,
+          model: nn.Module,
+          closure: Callable) -> Tuple[nn.Module]:
+        """
+        A model that is followed by a closure function.
+
+        Args:  
+        - `model`: The model in the chain.
+        - `closure`: Differentiable function to be called at the end of forward.
+
+        Returns:
+            Tuple[nn.Module]: A tuple containing the instantiated model.
+        """
+        with torch.inference_mode(False):
+            model = PtnModelWithClosureDef(model, closure)
+        return (model,)

modules/pytorch_wrapper/ptn_multihead_attention.py

@@ -0,0 +1,153 @@
+from typing import Any, Dict, Tuple
+import torch
+import torch.nn as nn
+from typing import Tuple
+
+class PtnMultiheadAttentionDef(nn.Module):
+    """
+    Implements MultiheadAttention wrapper.
+
+    Attributes:
+        embed_dim (int): Hidden dimension (d_model) for each token of the Transformer.
+        num_heads  (int): Number of attention heads.
+        dropout (float): Dropout rate for output.
+        bias (bool): If True, bias will be added to the linear projections of query, key, value, and the final output.
+        add_bias_kv (bool): If True, bias will be added to k and v sequences.
+        add_zero_attn (bool): If True, a zero batch will be added  to k and v sequences.
+        kdim (int): Dimension of k.
+        vdim (int): Dimension of v.
+        batch_first (bool): Input shape format: True for (B, Seq, Feature), False for (Seq, B, Feature).
+
+    pragma: skip_doc
+    """
+    def __init__(self,
+        embed_dim: int,
+        num_heads: int,
+        dropout: float,
+        bias: bool,
+        add_bias_kv: bool,
+        add_zero_attn: bool,
+        kdim: int,
+        vdim: int,
+        batch_first: bool
+    ) -> None:
+        super().__init__()
+
+        self.model = nn.MultiheadAttention(
+            embed_dim=embed_dim,
+            num_heads=num_heads,
+            dropout=dropout,
+            bias=bias,
+            add_bias_kv=add_bias_kv,
+            add_zero_attn=add_zero_attn,
+            kdim=kdim,
+            vdim=vdim,
+            batch_first=batch_first
+        )
+
+    def forward(self, inputs: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            inputs (torch.Tensor): Input tensor of shape (batch_size, seq_len, embed_dim) if batch_first is True.
+            mask (torch.Tensor): Attention mask where 1 indicates a valid token and 0 indicates padding.
+                This is the inverse of PyTorch's expected key_padding_mask, which masks out positions where the value is True.
+
+        Returns:
+            torch.Tensor: Output tensor after multi-head self-attention.
+        """
+        return self.model(
+            inputs, inputs, inputs,
+            need_weights=False,
+            key_padding_mask=~mask.to(torch.bool)
+        )[0]
+
+
+class PtnMultiheadAttention:
+    """
+    A Multihead attention model.
+
+        Args:
+            embed_dim (int): Hidden dimension (d_model) for each token of the Transformer.
+            num_heads  (int): Number of attention heads.
+            dropout (float): Dropout rate for output.
+            bias (bool): If True, bias will be added to the linear projections of query, key, value, and the final output.
+            add_bias_kv (bool): If True, bias will be added to k and v sequences.
+            add_zero_attn (bool): If True, a zero batch will be added  to k and v sequences.
+            kdim (int): Dimension of k.
+            vdim (int): Dimension of v.
+            batch_first (bool): Input shape format: True for (B, Seq, Feature), False for (Seq, B, Feature).
+            
+    category: PyTorch wrapper - Model
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types for the function.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required input types.
+        """
+        return {
+            "required": {
+                "embed_dim": ("INT", {"default": 512, "min": 1, "max": 1e6}),
+                "num_heads": ("INT", {"default": 8, "min": 1, "max": 256}),
+                "dropout": ("FLOAT", {"default": 0, "min": 0, "max": 1}),
+                "bias": ("BOOLEAN", {"default": True}),
+                "add_bias_kv": ("BOOLEAN", {"default": False}),
+                "add_zero_attn": ("BOOLEAN", {"default": False}),
+                "kdim": ("INT", {"default": 512, "min": 1, "max": 1e6}),
+                "vdim": ("INT", {"default": 512, "min": 1, "max": 1e6}),
+                "batch_first": ("BOOLEAN", {"default": True})
+            }
+        }
+
+    RETURN_TYPES: Tuple[str] = ("PTMODEL",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Training"
+
+    @classmethod
+    def IS_CHANGED(cls, **kw):
+        return float("NaN")
+
+    def f(self,
+        embed_dim: int,
+        num_heads: int,
+        dropout: float,
+        bias: bool,
+        add_bias_kv: bool,
+        add_zero_attn: bool,
+        kdim: int,
+        vdim: int,
+        batch_first: bool
+        ) -> Tuple[nn.Module]:
+        """
+        Constructs a Multihead attention model.
+
+        Args:
+            embed_dim (int): Hidden dimension (d_model) for each token of the Transformer.
+            num_heads  (int): Number of attention heads.
+            dropout (float): Dropout rate for output.
+            bias (bool): If True, bias will be added to the linear projections of query, key, value, and the final output.
+            add_bias_kv (bool): If True, a learnable bias vector is prepended to the key and value sequences along the sequence dimension. Each bias vector has the same dimension as other sequence positions. Set this to False if you are unsure.
+            add_zero_attn (bool): If True, a zero vector will be added to k and v sequences.
+            kdim (int): Dimension of k.
+            vdim (int): Dimension of v.
+            batch_first (bool): Input shape format: True for (B, Seq, Feature), False for (Seq, B, Feature).
+
+        Returns:
+            Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
+        """
+        with torch.inference_mode(False):
+            model = PtnMultiheadAttentionDef(
+                embed_dim=embed_dim,
+                num_heads=num_heads,
+                dropout=dropout,
+                bias=bias,
+                add_bias_kv=add_bias_kv,
+                add_zero_attn=add_zero_attn,
+                kdim=kdim,
+                vdim=vdim,
+                batch_first=batch_first
+            )
+            return (model,)

modules/pytorch_wrapper/ptn_residual_connection_model.py

@@ -0,0 +1,95 @@
+from typing import Any, Dict, Tuple, Callable, Optional
+import torch
+import torch.nn as nn
+
+
+class PtnResidualConnectionModelDef(nn.Module):
+    """
+    A model that saves the input and add to the output of the specified model.
+    optionally followed by a closure function.
+
+    ### Fields:  
+    - `model`: A model to process input.
+    - `closure`: Optional differentiable function to be called at the end of forward.
+
+    pragma: skip_doc
+    """
+
+    def __init__(self,
+                 model: nn.Module,
+                 closure: Optional[Callable] = None) -> None:
+        super().__init__()
+        self.model = model
+        self.closure = closure
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass through the chained models.
+
+        Args:
+            inputs (torch.Tensor): Input tensor.
+
+        Returns:
+            torch.Tensor: Output tensor after passing through `model` and adding residual input followed optionally by `closure`.
+        """
+        res = inputs
+        x = self.model(inputs)
+        x = res + x
+        if self.closure:
+            x = self.closure(x)
+        return x
+
+
+class PtnResidualConnectionModel:
+    """
+    A model that saves the input and add to the output of the specified model.
+    optionally followed by a closure function.
+
+    ### Fields:  
+    - `model`: A model to process input.
+    - `closure`: Optional differentiable function to be called at the end of forward.
+
+    category: PyTorch wrapper - Model
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the expected input types.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required and optional input types.
+        """
+        return {
+            "required": {
+                "model": ("PTMODEL", {})
+            },
+            "optional": {
+                "closure": ("PTCALLABLE", {})
+            }
+        }
+
+    RETURN_TYPES: Tuple[str, ...] = ("PTMODEL",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Training"
+
+    @classmethod
+    def IS_CHANGED(cls, **kw):
+        return float("NaN")
+
+    def f(self,
+          model: nn.Module,
+          closure: Optional[Callable] = None) -> Tuple[nn.Module]:
+        """
+        Constructs a chained model with optional closure.
+
+        Args:
+            model (nn.Module): A model that saves the input and add to the output of the specified model.
+            closure ( Optional[Callable]): An optional differentiable function to be called at the end of forward.
+
+        Returns:
+            Tuple[nn.Module]: A tuple containing the instantiated model.
+        """
+        with torch.inference_mode(False):
+            model = PtnResidualConnectionModelDef(model, closure)
+        return (model,)

modules/pytorch_wrapper/ptn_residual_connection_model_with_attention_mask.py

@@ -0,0 +1,96 @@
+from typing import Any, Dict, Tuple, Callable, Optional
+import torch
+import torch.nn as nn
+
+
+class PtnResidualConnectionModelWithAttentionMaskDef(nn.Module):
+    """
+    A model that saves the input and add to the output of the specified model.
+    optionally followed by a closure function.
+
+    ### Fields:  
+    - `model`: A model to process input.
+    - `closure`: Optional differentiable function to be called at the end of forward.
+
+    pragma: skip_doc
+    """
+
+    def __init__(self,
+                 model: nn.Module,
+                 closure: Optional[Callable] = None) -> None:
+        super().__init__()
+        self.model = model
+        self.closure = closure
+
+    def forward(self, inputs: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass through the chained models.
+
+        Args:
+            inputs (torch.Tensor): Input tensor.
+            mask (torch.Tensor): Attention mask.
+
+        Returns:
+            torch.Tensor: Output tensor after passing through `model` and adding residual input followed optionally by `closure`.
+        """
+        res = inputs
+        x = self.model(inputs, mask)
+        x = res + x
+        if self.closure:
+            x = self.closure(x)
+        return x
+
+
+class PtnResidualConnectionModelWithAttentionMask:
+    """
+    A model that saves the input and add to the output of the specified model.
+    optionally followed by a closure function.
+
+    ### Fields:  
+    - `model`: A model to process input.
+    - `closure`: Optional differentiable function to be called at the end of forward.
+
+    category: PyTorch wrapper - Model
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the expected input types.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required and optional input types.
+        """
+        return {
+            "required": {
+                "model": ("PTMODEL", {})
+            },
+            "optional": {
+                "closure": ("PTCALLABLE", {})
+            }
+        }
+
+    RETURN_TYPES: Tuple[str, ...] = ("PTMODEL",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Training"
+
+    @classmethod
+    def IS_CHANGED(cls, **kw):
+        return float("NaN")
+
+    def f(self,
+          model: nn.Module,
+          closure: Optional[Callable] = None) -> Tuple[nn.Module]:
+        """
+        Constructs a chained model with optional closure.
+
+        Args:
+            model (nn.Module): A model that saves the input and add to the output of the specified model.
+            closure ( Optional[Callable]): An optional differentiable function to be called at the end of forward.
+
+        Returns:
+            Tuple[nn.Module]: A tuple containing the instantiated model.
+        """
+        with torch.inference_mode(False):
+            model = PtnResidualConnectionModelWithAttentionMaskDef(model, closure)
+        return (model,)