diff --git a/README.md b/README.md
index 53f43c0f3..64a2ed0d5 100644
--- a/README.md
+++ b/README.md
@@ -90,20 +90,21 @@ Do you want to learn more about it? Look at our [Tutorials](https://github.com/m
 ### Solve Data Driven Problems
 Data driven modelling aims to learn a function that given some input data gives an output (e.g. regression, classification, ...). In PINA you can easily do this by:
 ```python
+import torch
 from pina import Trainer
 from pina.model import FeedForward
 from pina.solver import SupervisedSolver
 from pina.problem.zoo import SupervisedProblem
 
 input_tensor  = torch.rand((10, 1))
-output_tensor = input_tensor.pow(3)
+target_tensor = input_tensor.pow(3)
 
 # Step 1. Define problem
 problem = SupervisedProblem(input_tensor, target_tensor)
 # Step 2. Design model (you can use your favourite torch.nn.Module in here)
 model   = FeedForward(input_dimensions=1, output_dimensions=1, layers=[64, 64])
 # Step 3. Define Solver
-solver  = SupervisedSolver(problem, model)
+solver  = SupervisedSolver(problem, model, use_lt=False)
 # Step 4. Train
 trainer = Trainer(solver, max_epochs=1000, accelerator='gpu')
 trainer.train()
@@ -149,6 +150,7 @@ class SimpleODE(SpatialProblem):
 
 # Step 1. Define problem
 problem = SimpleODE()
+problem.discretise_domain(n=100, mode="grid", domains=["D", "x0"])
 # Step 2. Design model (you can use your favourite torch.nn.Module in here)
 model   = FeedForward(input_dimensions=1, output_dimensions=1, layers=[64, 64])
 # Step 3. Define Solver
diff --git a/docs/source/_rst/_code.rst b/docs/source/_rst/_code.rst
index 365537943..9bd36ab2d 100644
--- a/docs/source/_rst/_code.rst
+++ b/docs/source/_rst/_code.rst
@@ -104,6 +104,7 @@ Models
     LowRankNeuralOperator <model/low_rank_neural_operator.rst>
     GraphNeuralOperator <model/graph_neural_operator.rst>
     GraphNeuralKernel <model/graph_neural_operator_integral_kernel.rst>
+    PirateNet <model/pirate_network.rst>
 
 Blocks
 -------------
@@ -121,6 +122,7 @@ Blocks
     Continuous Convolution Interface <model/block/convolution_interface.rst>
     Continuous Convolution Block <model/block/convolution.rst>
     Orthogonal Block <model/block/orthogonal.rst>
+    PirateNet Block <model/block/pirate_network_block.rst>
 
 Message Passing
 -------------------
diff --git a/docs/source/_rst/model/block/pirate_network_block.rst b/docs/source/_rst/model/block/pirate_network_block.rst
new file mode 100644
index 000000000..5d0428a68
--- /dev/null
+++ b/docs/source/_rst/model/block/pirate_network_block.rst
@@ -0,0 +1,8 @@
+PirateNet Block
+=======================================
+.. currentmodule:: pina.model.block.pirate_network_block
+    
+.. autoclass:: PirateNetBlock
+    :members:
+    :show-inheritance:
+
diff --git a/docs/source/_rst/model/pirate_network.rst b/docs/source/_rst/model/pirate_network.rst
new file mode 100644
index 000000000..5b374c247
--- /dev/null
+++ b/docs/source/_rst/model/pirate_network.rst
@@ -0,0 +1,7 @@
+PirateNet
+=======================
+.. currentmodule:: pina.model.pirate_network
+
+.. autoclass:: PirateNet
+   :members:
+   :show-inheritance:
\ No newline at end of file
diff --git a/pina/callback/optimizer_callback.py b/pina/callback/optimizer_callback.py
index fb2770a43..1b518406b 100644
--- a/pina/callback/optimizer_callback.py
+++ b/pina/callback/optimizer_callback.py
@@ -21,26 +21,30 @@ def __init__(self, new_optimizers, epoch_switch):
             single :class:`torch.optim.Optimizer` instance or a list of them
             for multiple model solver.
         :type new_optimizers: pina.optim.TorchOptimizer | list
-        :param epoch_switch: The epoch at which the optimizer switch occurs.
-        :type epoch_switch: int
+        :param int epoch_switch: The epoch at which the optimizer switch occurs.
 
         Example:
-            >>> switch_callback = SwitchOptimizer(new_optimizers=optimizer,
-            >>>                                  epoch_switch=10)
+            >>> optimizer = TorchOptimizer(torch.optim.Adam, lr=0.01)
+            >>> switch_callback = SwitchOptimizer(
+            >>>     new_optimizers=optimizer, epoch_switch=10
+            >>> )
         """
         super().__init__()
 
+        # Check if epoch_switch is greater than 1
         if epoch_switch < 1:
             raise ValueError("epoch_switch must be greater than one.")
 
+        # If new_optimizers is not a list, convert it to a list
         if not isinstance(new_optimizers, list):
             new_optimizers = [new_optimizers]
 
-        # check type consistency
+        # Check consistency
+        check_consistency(epoch_switch, int)
         for optimizer in new_optimizers:
             check_consistency(optimizer, TorchOptimizer)
-        check_consistency(epoch_switch, int)
-        # save new optimizers
+
+        # Store the new optimizers and epoch switch
         self._new_optimizers = new_optimizers
         self._epoch_switch = epoch_switch
 
@@ -48,18 +52,21 @@ def on_train_epoch_start(self, trainer, __):
         """
         Switch the optimizer at the start of the specified training epoch.
 
-        :param trainer: The trainer object managing the training process.
-        :type trainer: pytorch_lightning.Trainer
+        :param lightning.pytorch.Trainer trainer: The trainer object managing
+            the training process.
         :param _: Placeholder argument (not used).
-
-        :return: None
-        :rtype: None
         """
+        # Check if the current epoch matches the switch epoch
         if trainer.current_epoch == self._epoch_switch:
             optims = []
 
+            # Hook the new optimizers to the model parameters
             for idx, optim in enumerate(self._new_optimizers):
                 optim.hook(trainer.solver._pina_models[idx].parameters())
                 optims.append(optim)
 
+            # Update the solver's optimizers
             trainer.solver._pina_optimizers = optims
+
+            # Update the trainer's strategy optimizers
+            trainer.strategy.optimizers = [o.instance for o in optims]
diff --git a/pina/equation/system_equation.py b/pina/equation/system_equation.py
index d51ba9408..21cb27160 100644
--- a/pina/equation/system_equation.py
+++ b/pina/equation/system_equation.py
@@ -8,18 +8,51 @@
 
 class SystemEquation(EquationInterface):
     """
-    Implementation of the System of Equations. Every ``equation`` passed to a
-    :class:`~pina.condition.condition.Condition` object must be either a
-    :class:`~pina.equation.equation.Equation` or a
-    :class:`~pina.equation.system_equation.SystemEquation` instance.
+    Implementation of the System of Equations, to be passed to a
+    :class:`~pina.condition.condition.Condition` object.
+
+    Unlike the :class:`~pina.equation.equation.Equation` class, which represents
+    a single equation, the :class:`SystemEquation` class allows multiple
+    equations to be grouped together into a system. This is particularly useful
+    when dealing with multi-component outputs or coupled physical models, where
+    the residual must be computed collectively across several constraints.
+
+    Each equation in the system must be either:
+    - An instance of :class:`~pina.equation.equation.Equation`;
+    - A callable function.
+
+    The residuals from each equation are computed independently and then
+    aggregated using an optional reduction strategy (e.g., ``mean``, ``sum``).
+    The resulting residual is returned as a single :class:`~pina.LabelTensor`.
+
+    :Example:
+
+    >>> from pina.equation import SystemEquation, FixedValue, FixedGradient
+    >>> from pina import LabelTensor
+    >>> import torch
+    >>> pts = LabelTensor(torch.rand(10, 2), labels=["x", "y"])
+    >>> pts.requires_grad = True
+    >>> output_ = torch.pow(pts, 2)
+    >>> output_.labels = ["u", "v"]
+    >>> system_equation = SystemEquation(
+    ...     [
+    ...         FixedValue(value=1.0, components=["u"]),
+    ...         FixedGradient(value=0.0, components=["v"],d=["y"]),
+    ...     ],
+    ...     reduction="mean",
+    ... )
+    >>> residual = system_equation.residual(pts, output_)
+
     """
 
     def __init__(self, list_equation, reduction=None):
         """
         Initialization of the :class:`SystemEquation` class.
 
-        :param Callable equation: A ``torch`` callable function used to compute
-            the residual of a mathematical equation.
+        :param list_equation: A list containing either callable functions or
+            instances of :class:`~pina.equation.equation.Equation`, used to
+            compute the residuals of mathematical equations.
+        :type list_equation: list[Callable] | list[Equation]
         :param str reduction: The reduction method to aggregate the residuals of
             each equation. Available options are: ``None``, ``mean``, ``sum``,
             ``callable``.
@@ -32,9 +65,10 @@ def __init__(self, list_equation, reduction=None):
         check_consistency([list_equation], list)
 
         # equations definition
-        self.equations = []
-        for _, equation in enumerate(list_equation):
-            self.equations.append(Equation(equation))
+        self.equations = [
+            equation if isinstance(equation, Equation) else Equation(equation)
+            for equation in list_equation
+        ]
 
         # possible reduction
         if reduction == "mean":
@@ -45,7 +79,7 @@ def __init__(self, list_equation, reduction=None):
             self.reduction = reduction
         else:
             raise NotImplementedError(
-                "Only mean and sum reductions implemented."
+                "Only mean and sum reductions are currenly supported."
             )
 
     def residual(self, input_, output_, params_=None):
diff --git a/pina/model/__init__.py b/pina/model/__init__.py
index 606dde7d5..5e340484e 100644
--- a/pina/model/__init__.py
+++ b/pina/model/__init__.py
@@ -13,6 +13,7 @@
     "LowRankNeuralOperator",
     "Spline",
     "GraphNeuralOperator",
+    "PirateNet",
 ]
 
 from .feed_forward import FeedForward, ResidualFeedForward
@@ -24,3 +25,4 @@
 from .low_rank_neural_operator import LowRankNeuralOperator
 from .spline import Spline
 from .graph_neural_operator import GraphNeuralOperator
+from .pirate_network import PirateNet
diff --git a/pina/model/block/__init__.py b/pina/model/block/__init__.py
index c40135b7e..08b313387 100644
--- a/pina/model/block/__init__.py
+++ b/pina/model/block/__init__.py
@@ -18,6 +18,7 @@
     "LowRankBlock",
     "RBFBlock",
     "GNOBlock",
+    "PirateNetBlock",
 ]
 
 from .convolution_2d import ContinuousConvBlock
@@ -35,3 +36,4 @@
 from .low_rank_block import LowRankBlock
 from .rbf_block import RBFBlock
 from .gno_block import GNOBlock
+from .pirate_network_block import PirateNetBlock
diff --git a/pina/model/block/pirate_network_block.py b/pina/model/block/pirate_network_block.py
new file mode 100644
index 000000000..cfeb8410e
--- /dev/null
+++ b/pina/model/block/pirate_network_block.py
@@ -0,0 +1,89 @@
+"""Module for the PirateNet block class."""
+
+import torch
+from ...utils import check_consistency, check_positive_integer
+
+
+class PirateNetBlock(torch.nn.Module):
+    """
+    The inner block of Physics-Informed residual adaptive network (PirateNet).
+
+    The block consists of three dense layers with dual gating operations and an
+    adaptive residual connection. The trainable ``alpha`` parameter controls
+    the contribution of the residual connection.
+
+    .. seealso::
+
+        **Original reference**:
+        Wang, S., Sankaran, S., Stinis., P., Perdikaris, P. (2025).
+        *Simulating Three-dimensional Turbulence with Physics-informed Neural
+        Networks*.
+        DOI: `arXiv preprint arXiv:2507.08972.
+        <https://arxiv.org/abs/2507.08972>`_
+    """
+
+    def __init__(self, inner_size, activation):
+        """
+        Initialization of the :class:`PirateNetBlock` class.
+
+        :param int inner_size: The number of hidden units in the dense layers.
+        :param torch.nn.Module activation: The activation function.
+        """
+        super().__init__()
+
+        # Check consistency
+        check_consistency(activation, torch.nn.Module, subclass=True)
+        check_positive_integer(inner_size, strict=True)
+
+        # Initialize the linear transformations of the dense layers
+        self.linear1 = torch.nn.Linear(inner_size, inner_size)
+        self.linear2 = torch.nn.Linear(inner_size, inner_size)
+        self.linear3 = torch.nn.Linear(inner_size, inner_size)
+
+        # Initialize the scales of the dense layers
+        self.scale1 = torch.nn.Parameter(torch.zeros(inner_size))
+        self.scale2 = torch.nn.Parameter(torch.zeros(inner_size))
+        self.scale3 = torch.nn.Parameter(torch.zeros(inner_size))
+
+        # Initialize the adaptive residual connection parameter
+        self._alpha = torch.nn.Parameter(torch.zeros(1))
+
+        # Initialize the activation function
+        self.activation = activation()
+
+    def forward(self, x, U, V):
+        """
+        Forward pass of the PirateNet block. It computes the output of the block
+        by applying the dense layers with scaling, and combines the results with
+        the input using the adaptive residual connection.
+
+        :param x: The input tensor.
+        :type x: torch.Tensor | LabelTensor
+        :param torch.Tensor U: The first shared gating tensor. It must have the
+            same shape as ``x``.
+        :param torch.Tensor V: The second shared gating tensor. It must have the
+            same shape as ``x``.
+        :return: The output tensor of the block.
+        :rtype: torch.Tensor | LabelTensor
+        """
+        # Compute the output of the first dense layer with scaling
+        f = self.activation(self.linear1(x) * torch.exp(self.scale1))
+        z1 = f * U + (1 - f) * V
+
+        # Compute the output of the second dense layer with scaling
+        g = self.activation(self.linear2(z1) * torch.exp(self.scale2))
+        z2 = g * U + (1 - g) * V
+
+        # Compute the output of the block
+        h = self.activation(self.linear3(z2) * torch.exp(self.scale3))
+        return self._alpha * h + (1 - self._alpha) * x
+
+    @property
+    def alpha(self):
+        """
+        Return the alpha parameter.
+
+        :return: The alpha parameter controlling the residual connection.
+        :rtype: torch.nn.Parameter
+        """
+        return self._alpha
diff --git a/pina/model/pirate_network.py b/pina/model/pirate_network.py
new file mode 100644
index 000000000..96102b41f
--- /dev/null
+++ b/pina/model/pirate_network.py
@@ -0,0 +1,118 @@
+"""Module for the PirateNet model class."""
+
+import torch
+from .block import FourierFeatureEmbedding, PirateNetBlock
+from ..utils import check_consistency, check_positive_integer
+
+
+class PirateNet(torch.nn.Module):
+    """
+    Implementation of Physics-Informed residual adaptive network (PirateNet).
+
+    The model consists of a Fourier feature embedding layer, multiple PirateNet
+    blocks, and a final output layer. Each PirateNet block consist of three
+    dense layers with dual gating mechanism and an adaptive residual connection,
+    whose contribution is controlled by a trainable parameter ``alpha``.
+
+    The PirateNet, augmented with random weight factorization, is designed to
+    mitigate spectral bias in deep networks.
+
+    .. seealso::
+
+        **Original reference**:
+        Wang, S., Sankaran, S., Stinis., P., Perdikaris, P. (2025).
+        *Simulating Three-dimensional Turbulence with Physics-informed Neural
+        Networks*.
+        DOI: `arXiv preprint arXiv:2507.08972.
+        <https://arxiv.org/abs/2507.08972>`_
+    """
+
+    def __init__(
+        self,
+        input_dimension,
+        inner_size,
+        output_dimension,
+        embedding=None,
+        n_layers=3,
+        activation=torch.nn.Tanh,
+    ):
+        """
+        Initialization of the :class:`PirateNet` class.
+
+        :param int input_dimension: The number of input features.
+        :param int inner_size: The number of hidden units in the dense layers.
+        :param int output_dimension: The number of output features.
+        :param torch.nn.Module embedding: The embedding module used to transform
+            the input into a higher-dimensional feature space. If ``None``, a
+            default :class:`~pina.model.block.FourierFeatureEmbedding` with
+            scaling factor of 2 is used. Default is ``None``.
+        :param int n_layers: The number of PirateNet blocks in the model.
+            Default is 3.
+        :param torch.nn.Module activation: The activation function to be used in
+            the blocks. Default is :class:`torch.nn.Tanh`.
+        """
+        super().__init__()
+
+        # Check consistency
+        check_consistency(activation, torch.nn.Module, subclass=True)
+        check_positive_integer(input_dimension, strict=True)
+        check_positive_integer(inner_size, strict=True)
+        check_positive_integer(output_dimension, strict=True)
+        check_positive_integer(n_layers, strict=True)
+
+        # Initialize the activation function
+        self.activation = activation()
+
+        # Initialize the Fourier embedding
+        self.embedding = embedding or FourierFeatureEmbedding(
+            input_dimension=input_dimension,
+            output_dimension=inner_size,
+            sigma=2.0,
+        )
+
+        # Initialize the shared dense layers
+        self.linear1 = torch.nn.Linear(inner_size, inner_size)
+        self.linear2 = torch.nn.Linear(inner_size, inner_size)
+
+        # Initialize the PirateNet blocks
+        self.blocks = torch.nn.ModuleList(
+            [PirateNetBlock(inner_size, activation) for _ in range(n_layers)]
+        )
+
+        # Initialize the output layer
+        self.output_layer = torch.nn.Linear(inner_size, output_dimension)
+
+    def forward(self, input_):
+        """
+        Forward pass of the PirateNet model. It applies the Fourier feature
+        embedding, computes the shared gating tensors U and V, and passes the
+        input through each block in the network. Finally, it applies the output
+        layer to produce the final output.
+
+        :param input_: The input tensor for the model.
+        :type input_: torch.Tensor | LabelTensor
+        :return: The output tensor of the model.
+        :rtype: torch.Tensor | LabelTensor
+        """
+        # Apply the Fourier feature embedding
+        x = self.embedding(input_)
+
+        # Compute U and V from the shared dense layers
+        U = self.activation(self.linear1(x))
+        V = self.activation(self.linear2(x))
+
+        # Pass through each block in the network
+        for block in self.blocks:
+            x = block(x, U, V)
+
+        return self.output_layer(x)
+
+    @property
+    def alpha(self):
+        """
+        Return the alpha values of all PirateNetBlock layers.
+
+        :return: A list of alpha values from each block.
+        :rtype: list
+        """
+        return [block.alpha.item() for block in self.blocks]
diff --git a/pina/problem/zoo/inverse_poisson_2d_square.py b/pina/problem/zoo/inverse_poisson_2d_square.py
index f112ebfc0..ab89c4c16 100644
--- a/pina/problem/zoo/inverse_poisson_2d_square.py
+++ b/pina/problem/zoo/inverse_poisson_2d_square.py
@@ -1,5 +1,6 @@
 """Formulation of the inverse Poisson problem in a square domain."""
 
+import warnings
 import requests
 import torch
 from io import BytesIO
@@ -9,6 +10,44 @@
 from ...domain import CartesianDomain
 from ...equation import Equation, FixedValue
 from ...problem import SpatialProblem, InverseProblem
+from ...utils import custom_warning_format, check_consistency
+
+warnings.formatwarning = custom_warning_format
+warnings.filterwarnings("always", category=ResourceWarning)
+
+
+def _load_tensor_from_url(url, labels, timeout=10):
+    """
+    Downloads a tensor file from a URL and wraps it in a LabelTensor.
+
+    This function fetches a `.pth` file containing tensor data, extracts it,
+    and returns it as a LabelTensor using the specified labels. If the file
+    cannot be retrieved (e.g., no internet connection), a warning is issued
+    and None is returned.
+
+    :param str url: URL to the remote `.pth` tensor file.
+    :param list[str] | tuple[str] labels: Labels for the resulting LabelTensor.
+    :param int timeout: Timeout for the request in seconds.
+    :return: A LabelTensor object if successful, otherwise None.
+    :rtype: LabelTensor | None
+    """
+    # Try to download the tensor file from the given URL
+    try:
+        response = requests.get(url, timeout=timeout)
+        response.raise_for_status()
+        tensor = torch.load(
+            BytesIO(response.content), weights_only=False
+        ).tensor.detach()
+        return LabelTensor(tensor, labels)
+
+    # If the request fails, issue a warning and return None
+    except requests.exceptions.RequestException as e:
+        warnings.warn(
+            f"Could not download data for 'InversePoisson2DSquareProblem' "
+            f"from '{url}'. Reason: {e}. Skipping data loading.",
+            ResourceWarning,
+        )
+        return None
 
 
 def laplace_equation(input_, output_, params_):
@@ -29,35 +68,13 @@ def laplace_equation(input_, output_, params_):
     return delta_u - force_term
 
 
-# URL of the file
-url = "https://github.com/mathLab/PINA/raw/refs/heads/master/tutorials/tutorial7/data/pts_0.5_0.5"
-# Download the file
-response = requests.get(url)
-response.raise_for_status()
-file_like_object = BytesIO(response.content)
-# Set the data
-input_data = LabelTensor(
-    torch.load(file_like_object, weights_only=False).tensor.detach(),
-    ["x", "y", "mu1", "mu2"],
-)
-
-# URL of the file
-url = "https://github.com/mathLab/PINA/raw/refs/heads/master/tutorials/tutorial7/data/pinn_solution_0.5_0.5"
-# Download the file
-response = requests.get(url)
-response.raise_for_status()
-file_like_object = BytesIO(response.content)
-# Set the data
-output_data = LabelTensor(
-    torch.load(file_like_object, weights_only=False).tensor.detach(), ["u"]
-)
-
-
 class InversePoisson2DSquareProblem(SpatialProblem, InverseProblem):
     r"""
     Implementation of the inverse 2-dimensional Poisson problem in the square
     domain :math:`[0, 1] \times [0, 1]`,
     with unknown parameter domain :math:`[-1, 1] \times [-1, 1]`.
+    The `"data"` condition is added only if the required files are
+    downloaded successfully.
 
     :Example:
         >>> problem = InversePoisson2DSquareProblem()
@@ -83,5 +100,52 @@ class InversePoisson2DSquareProblem(SpatialProblem, InverseProblem):
         "g3": Condition(domain="g3", equation=FixedValue(0.0)),
         "g4": Condition(domain="g4", equation=FixedValue(0.0)),
         "D": Condition(domain="D", equation=Equation(laplace_equation)),
-        "data": Condition(input=input_data, target=output_data),
     }
+
+    def __init__(self, load=True, data_size=1.0):
+        """
+        Initialization of the :class:`InversePoisson2DSquareProblem`.
+
+        :param bool load: If True, it attempts to load data from remote URLs.
+            Set to False to skip data loading (e.g., if no internet connection).
+        :param float data_size: The fraction of the total data to use for the
+            "data" condition. If set to 1.0, all available data is used.
+            If set to 0.0, no data is used. Default is 1.0.
+        :raises ValueError: If `data_size` is not in the range [0.0, 1.0].
+        :raises ValueError: If `data_size` is not a float.
+        """
+        super().__init__()
+
+        # Check consistency
+        check_consistency(load, bool)
+        check_consistency(data_size, float)
+        if not 0.0 <= data_size <= 1.0:
+            raise ValueError(
+                f"data_size must be in the range [0.0, 1.0], got {data_size}."
+            )
+
+        # Load data if requested
+        if load:
+
+            # Define URLs for input and output data
+            input_url = (
+                "https://github.com/mathLab/PINA/raw/refs/heads/master"
+                "/tutorials/tutorial7/data/pts_0.5_0.5"
+            )
+            output_url = (
+                "https://github.com/mathLab/PINA/raw/refs/heads/master"
+                "/tutorials/tutorial7/data/pinn_solution_0.5_0.5"
+            )
+
+            # Define input and output data
+            input_data = _load_tensor_from_url(
+                input_url, ["x", "y", "mu1", "mu2"]
+            )
+            output_data = _load_tensor_from_url(output_url, ["u"])
+
+            # Add the "data" condition
+            if input_data is not None and output_data is not None:
+                n_data = int(input_data.shape[0] * data_size)
+                self.conditions["data"] = Condition(
+                    input=input_data[:n_data], target=output_data[:n_data]
+                )
diff --git a/pina/solver/physics_informed_solver/rba_pinn.py b/pina/solver/physics_informed_solver/rba_pinn.py
index 808ac5a1b..5c8d50fed 100644
--- a/pina/solver/physics_informed_solver/rba_pinn.py
+++ b/pina/solver/physics_informed_solver/rba_pinn.py
@@ -1,6 +1,5 @@
 """Module for the Residual-Based Attention PINN solver."""
 
-from copy import deepcopy
 import torch
 
 from .pinn import PINN
@@ -98,6 +97,8 @@ def __init__(
         :param float gamma: The decay parameter in the update of the weights
             of the residuals. Must be between ``0`` and ``1``.
             Default is ``0.999``.
+        :raises: ValueError if `gamma` is not in the range (0, 1).
+        :raises: ValueError if `eta` is not greater than 0.
         """
         super().__init__(
             model=model,
@@ -111,78 +112,216 @@ def __init__(
         # check consistency
         check_consistency(eta, (float, int))
         check_consistency(gamma, float)
-        assert (
-            0 < gamma < 1
-        ), f"Invalid range: expected 0 < gamma < 1, got {gamma=}"
+
+        # Validate range for gamma
+        if not 0 < gamma < 1:
+            raise ValueError(
+                f"Invalid range: expected 0 < gamma < 1, but got {gamma}"
+            )
+
+        # Validate range for eta
+        if eta <= 0:
+            raise ValueError(f"Invalid range: expected eta > 0, but got {eta}")
+
+        # Initialize parameters
         self.eta = eta
         self.gamma = gamma
 
-        # initialize weights
+        # Initialize the weight of each point to 0
         self.weights = {}
-        for condition_name in problem.conditions:
-            self.weights[condition_name] = 0
+        for cond, data in self.problem.input_pts.items():
+            buffer_tensor = torch.zeros((len(data), 1), device=self.device)
+            self.register_buffer(f"weight_{cond}", buffer_tensor)
+            self.weights[cond] = getattr(self, f"weight_{cond}")
+
+        # Extract the reduction method from the loss function
+        self._reduction = self._loss_fn.reduction
 
-        # define vectorial loss
-        self._vectorial_loss = deepcopy(self.loss)
-        self._vectorial_loss.reduction = "none"
+        # Set the loss function to return non-aggregated losses
+        self._loss_fn = type(self._loss_fn)(reduction="none")
 
-    # for now RBAPINN is implemented only for batch_size = None
     def on_train_start(self):
         """
-        Hook method called at the beginning of training.
+        Ensure that all residual weight buffers registered during initialization
+        are moved to the correct computation device.
+        """
+        # Move all weight buffers to the correct device
+        for cond in self.problem.input_pts:
 
-        :raises NotImplementedError: If the batch size is not ``None``.
+            # Get the buffer for the current condition
+            weight_buf = getattr(self, f"weight_{cond}")
+
+            # Move the buffer to the correct device
+            weight_buf.data = weight_buf.data.to(self.device)
+            self.weights[cond] = weight_buf
+
+    def training_step(self, batch, batch_idx, **kwargs):
         """
-        if self.trainer.batch_size is not None:
-            raise NotImplementedError(
-                "RBAPINN only works with full batch "
-                "size, set batch_size=None inside the "
-                "Trainer to use the solver."
-            )
-        return super().on_train_start()
+        Solver training step. It computes the optimization cycle and aggregates
+        the losses using the ``weighting`` attribute.
 
-    def _vect_to_scalar(self, loss_value):
+        :param list[tuple[str, dict]] batch: A batch of data. Each element is a
+            tuple containing a condition name and a dictionary of points.
+        :param int batch_idx: The index of the current batch.
+        :param dict kwargs: Additional keyword arguments passed to
+            ``optimization_cycle``.
+        :return: The loss of the training step.
+        :rtype: torch.Tensor
         """
-        Computation of the scalar loss.
+        loss = self._optimization_cycle(
+            batch=batch, batch_idx=batch_idx, **kwargs
+        )
+        self.store_log("train_loss", loss, self.get_batch_size(batch))
+        return loss
 
-        :param LabelTensor loss_value: the tensor of pointwise losses.
-        :raises RuntimeError: If the loss reduction is not ``mean`` or ``sum``.
-        :return: The computed scalar loss.
-        :rtype: LabelTensor
+    @torch.set_grad_enabled(True)
+    def validation_step(self, batch, **kwargs):
         """
-        if self.loss.reduction == "mean":
-            ret = torch.mean(loss_value)
-        elif self.loss.reduction == "sum":
-            ret = torch.sum(loss_value)
-        else:
-            raise RuntimeError(
-                f"Invalid reduction, got {self.loss.reduction} "
-                "but expected mean or sum."
+        The validation step for the PINN solver. It returns the average residual
+        computed with the ``loss`` function not aggregated.
+
+        :param list[tuple[str, dict]] batch: A batch of data. Each element is a
+            tuple containing a condition name and a dictionary of points.
+        :param dict kwargs: Additional keyword arguments passed to
+            ``optimization_cycle``.
+        :return: The loss of the validation step.
+        :rtype: torch.Tensor
+        """
+        losses = self.optimization_cycle(batch=batch, **kwargs)
+
+        # Aggregate losses for each condition
+        for cond, loss in losses.items():
+            losses[cond] = self._apply_reduction(loss=losses[cond])
+
+        loss = (sum(losses.values()) / len(losses)).as_subclass(torch.Tensor)
+        self.store_log("val_loss", loss, self.get_batch_size(batch))
+        return loss
+
+    @torch.set_grad_enabled(True)
+    def test_step(self, batch, **kwargs):
+        """
+        The test step for the PINN solver. It returns the average residual
+        computed with the ``loss`` function not aggregated.
+
+        :param list[tuple[str, dict]] batch: A batch of data. Each element is a
+            tuple containing a condition name and a dictionary of points.
+        :param dict kwargs: Additional keyword arguments passed to
+            ``optimization_cycle``.
+        :return: The loss of the test step.
+        :rtype: torch.Tensor
+        """
+        losses = self.optimization_cycle(batch=batch, **kwargs)
+
+        # Aggregate losses for each condition
+        for cond, loss in losses.items():
+            losses[cond] = self._apply_reduction(loss=losses[cond])
+
+        loss = (sum(losses.values()) / len(losses)).as_subclass(torch.Tensor)
+        self.store_log("test_loss", loss, self.get_batch_size(batch))
+        return loss
+
+    def _optimization_cycle(self, batch, batch_idx, **kwargs):
+        """
+        Aggregate the loss for each condition in the batch.
+
+        :param list[tuple[str, dict]] batch: A batch of data. Each element is a
+            tuple containing a condition name and a dictionary of points.
+        :param int batch_idx: The index of the current batch.
+        :param dict kwargs: Additional keyword arguments passed to
+            ``optimization_cycle``.
+        :return: The losses computed for all conditions in the batch, casted
+            to a subclass of :class:`torch.Tensor`. It should return a dict
+            containing the condition name and the associated scalar loss.
+        :rtype: dict
+        """
+        # compute non-aggregated residuals
+        residuals = self.optimization_cycle(batch)
+
+        # update weights based on residuals
+        self._update_weights(batch, batch_idx, residuals)
+
+        # compute losses
+        losses = {}
+        for cond, res in residuals.items():
+
+            # Get the correct indices for the weights. Modulus is used according
+            # to the number of points in the condition, as in the PinaDataset.
+            len_res = len(res)
+            idx = torch.arange(
+                batch_idx * len_res,
+                (batch_idx + 1) * len_res,
+                device=self.weights[cond].device,
+            ) % len(self.problem.input_pts[cond])
+
+            losses[cond] = self._apply_reduction(
+                loss=(res * self.weights[cond][idx])
             )
-        return ret
 
-    def loss_phys(self, samples, equation):
+            # store log
+            self.store_log(
+                f"{cond}_loss", losses[cond].item(), self.get_batch_size(batch)
+            )
+
+        # clamp unknown parameters in InverseProblem (if needed)
+        self._clamp_params()
+
+        # aggregate
+        loss = self.weighting.aggregate(losses).as_subclass(torch.Tensor)
+
+        return loss
+
+    def _update_weights(self, batch, batch_idx, residuals):
         """
-        Computes the physics loss for the physics-informed solver based on the
-        provided samples and equation.
+        Update weights based on residuals.
 
-        :param LabelTensor samples: The samples to evaluate the physics loss.
-        :param EquationInterface equation: The governing equation.
-        :return: The computed physics loss.
-        :rtype: LabelTensor
+        :param list[tuple[str, dict]] batch: A batch of data. Each element is a
+            tuple containing a condition name and a dictionary of points.
+        :param int batch_idx: The index of the current batch.
+        :param dict residuals: A dictionary containing the residuals for each
+            condition. The keys are the condition names and the values are the
+            residuals as tensors.
         """
-        residual = self.compute_residual(samples=samples, equation=equation)
-        cond = self.current_condition_name
+        # Iterate over each condition in the batch
+        for cond, data in batch:
 
-        r_norm = (
-            self.eta
-            * torch.abs(residual)
-            / (torch.max(torch.abs(residual)) + 1e-12)
-        )
-        self.weights[cond] = (self.gamma * self.weights[cond] + r_norm).detach()
+            # Compute normalized residuals
+            res = residuals[cond]
+            res_abs = torch.linalg.vector_norm(res, ord=2, dim=1, keepdim=True)
+            r_norm = (self.eta * res_abs) / (res_abs.max() + 1e-12)
 
-        loss_value = self._vectorial_loss(
-            torch.zeros_like(residual, requires_grad=True), residual
-        )
+            # Get the correct indices for the weights. Modulus is used according
+            # to the number of points in the condition, as in the PinaDataset.
+            len_pts = len(data["input"])
+            idx = torch.arange(
+                batch_idx * len_pts,
+                (batch_idx + 1) * len_pts,
+                device=self.weights[cond].device,
+            ) % len(self.problem.input_pts[cond])
+
+            # Update weights
+            weights = self.weights[cond]
+            update = self.gamma * weights[idx] + r_norm
+            weights[idx] = update.detach()
 
-        return self._vect_to_scalar(self.weights[cond] ** 2 * loss_value)
+    def _apply_reduction(self, loss):
+        """
+        Apply the specified reduction to the loss. The reduction is deferred
+        until the end of the optimization cycle to allow residual-based weights
+        to be applied to each point beforehand.
+
+        :param torch.Tensor loss: The loss tensor to be reduced.
+        :return: The reduced loss tensor.
+        :rtype: torch.Tensor
+        :raises ValueError: If the reduction method is neither "mean" nor "sum".
+        """
+        # Apply the specified reduction method
+        if self._reduction == "mean":
+            return loss.mean()
+        if self._reduction == "sum":
+            return loss.sum()
+
+        # Raise an error if the reduction method is not recognized
+        raise ValueError(
+            f"Unknown reduction: {self._reduction}."
+            " Supported reductions are 'mean' and 'sum'."
+        )
diff --git a/pina/solver/physics_informed_solver/self_adaptive_pinn.py b/pina/solver/physics_informed_solver/self_adaptive_pinn.py
index 9521556e4..b1d2a2cb4 100644
--- a/pina/solver/physics_informed_solver/self_adaptive_pinn.py
+++ b/pina/solver/physics_informed_solver/self_adaptive_pinn.py
@@ -15,15 +15,20 @@ class Weights(torch.nn.Module):
     :class:`SelfAdaptivePINN` solver.
     """
 
-    def __init__(self, func):
+    def __init__(self, func, num_points):
         """
         Initialization of the :class:`Weights` class.
 
         :param torch.nn.Module func: the mask model.
+        :param int num_points: the number of input points.
         """
         super().__init__()
+
+        # Check consistency
         check_consistency(func, torch.nn.Module)
-        self.sa_weights = torch.nn.Parameter(torch.Tensor())
+
+        # Initialize the weights as a learnable parameter
+        self.sa_weights = torch.nn.Parameter(torch.zeros(num_points, 1))
         self.func = func
 
     def forward(self):
@@ -140,17 +145,17 @@ def __init__(
             If ``None``, the :class:`torch.nn.MSELoss` loss is used.
             Default is `None`.
         """
-        # check consistency weitghs_function
+        # Check consistency
         check_consistency(weight_function, torch.nn.Module)
 
-        # create models for weights
-        weights_dict = {}
-        for condition_name in problem.conditions:
-            weights_dict[condition_name] = Weights(weight_function)
-        weights_dict = torch.nn.ModuleDict(weights_dict)
+        # Define a ModuleDict for the weights
+        weights = {}
+        for cond, data in problem.input_pts.items():
+            weights[cond] = Weights(func=weight_function, num_points=len(data))
+        weights = torch.nn.ModuleDict(weights)
 
         super().__init__(
-            models=[model, weights_dict],
+            models=[model, weights],
             problem=problem,
             optimizers=[optimizer_model, optimizer_weights],
             schedulers=[scheduler_model, scheduler_weights],
@@ -158,44 +163,133 @@ def __init__(
             loss=loss,
         )
 
-        self._vectorial_loss = deepcopy(self.loss)
-        self._vectorial_loss.reduction = "none"
+        # Extract the reduction method from the loss function
+        self._reduction = self._loss_fn.reduction
 
-    def forward(self, x):
-        """
-        Forward pass.
+        # Set the loss function to return non-aggregated losses
+        self._loss_fn = type(self._loss_fn)(reduction="none")
 
-        :param LabelTensor x: Input tensor.
-        :return: The output of the neural network.
-        :rtype: LabelTensor
+    def training_step(self, batch, batch_idx, **kwargs):
         """
-        return self.model(x)
-
-    def training_step(self, batch):
-        """
-        Solver training step, overridden to perform manual optimization.
+        Solver training step. It computes the optimization cycle and aggregates
+        the losses using the ``weighting`` attribute.
 
         :param list[tuple[str, dict]] batch: A batch of data. Each element is a
             tuple containing a condition name and a dictionary of points.
-        :return: The aggregated loss.
-        :rtype: LabelTensor
+        :param int batch_idx: The index of the current batch.
+        :param dict kwargs: Additional keyword arguments passed to
+            ``optimization_cycle``.
+        :return: The loss of the training step.
+        :rtype: torch.Tensor
         """
         # Weights optimization
         self.optimizer_weights.instance.zero_grad()
-        loss = super().training_step(batch)
+        loss = self._optimization_cycle(
+            batch=batch, batch_idx=batch_idx, **kwargs
+        )
         self.manual_backward(-loss)
         self.optimizer_weights.instance.step()
         self.scheduler_weights.instance.step()
 
         # Model optimization
         self.optimizer_model.instance.zero_grad()
-        loss = super().training_step(batch)
+        loss = self._optimization_cycle(
+            batch=batch, batch_idx=batch_idx, **kwargs
+        )
         self.manual_backward(loss)
         self.optimizer_model.instance.step()
         self.scheduler_model.instance.step()
 
+        # Log the loss
+        self.store_log("train_loss", loss, self.get_batch_size(batch))
+
+        return loss
+
+    @torch.set_grad_enabled(True)
+    def validation_step(self, batch, **kwargs):
+        """
+        The validation step for the Self-Adaptive PINN solver. It returns the
+        average residual computed with the ``loss`` function not aggregated.
+
+        :param list[tuple[str, dict]] batch: A batch of data. Each element is a
+            tuple containing a condition name and a dictionary of points.
+        :param dict kwargs: Additional keyword arguments passed to
+            ``optimization_cycle``.
+        :return: The loss of the validation step.
+        :rtype: torch.Tensor
+        """
+        losses = self.optimization_cycle(batch=batch, **kwargs)
+
+        # Aggregate losses for each condition
+        for cond, loss in losses.items():
+            losses[cond] = self._apply_reduction(loss=losses[cond])
+
+        loss = (sum(losses.values()) / len(losses)).as_subclass(torch.Tensor)
+        self.store_log("val_loss", loss, self.get_batch_size(batch))
         return loss
 
+    @torch.set_grad_enabled(True)
+    def test_step(self, batch, **kwargs):
+        """
+        The test step for the Self-Adaptive PINN solver. It returns the average
+        residual computed with the ``loss`` function not aggregated.
+
+        :param list[tuple[str, dict]] batch: A batch of data. Each element is a
+            tuple containing a condition name and a dictionary of points.
+        :param dict kwargs: Additional keyword arguments passed to
+            ``optimization_cycle``.
+        :return: The loss of the test step.
+        :rtype: torch.Tensor
+        """
+        losses = self.optimization_cycle(batch=batch, **kwargs)
+
+        # Aggregate losses for each condition
+        for cond, loss in losses.items():
+            losses[cond] = self._apply_reduction(loss=losses[cond])
+
+        loss = (sum(losses.values()) / len(losses)).as_subclass(torch.Tensor)
+        self.store_log("test_loss", loss, self.get_batch_size(batch))
+        return loss
+
+    def loss_phys(self, samples, equation):
+        """
+        Computes the physics loss for the physics-informed solver based on the
+        provided samples and equation.
+
+        :param LabelTensor samples: The samples to evaluate the physics loss.
+        :param EquationInterface equation: The governing equation.
+        :return: The computed physics loss.
+        :rtype: LabelTensor
+        """
+        residuals = self.compute_residual(samples, equation)
+        return self._loss_fn(residuals, torch.zeros_like(residuals))
+
+    def loss_data(self, input, target):
+        """
+        Compute the data loss for the Self-Adaptive PINN solver by evaluating
+        the loss between the network's output and the true solution. This method
+        should not be overridden, if not intentionally.
+
+        :param input: The input to the neural network.
+        :type input: LabelTensor | torch.Tensor
+        :param target: The target to compare with the network's output.
+        :type target: LabelTensor | torch.Tensor
+        :return: The supervised loss, averaged over the number of observations.
+        :rtype: LabelTensor | torch.Tensor
+        """
+        return self._loss_fn(self.forward(input), target)
+
+    def forward(self, x):
+        """
+        Forward pass.
+
+        :param x: Input tensor.
+        :type x: torch.Tensor | LabelTensor
+        :return: The output of the neural network.
+        :rtype: torch.Tensor | LabelTensor
+        """
+        return self.model(x)
+
     def configure_optimizers(self):
         """
         Optimizer configuration.
@@ -203,10 +297,11 @@ def configure_optimizers(self):
         :return: The optimizers and the schedulers
         :rtype: tuple[list[Optimizer], list[Scheduler]]
         """
-        # If the problem is an InverseProblem, add the unknown parameters
-        # to the parameters to be optimized
+        # Hook the optimizers to the models
         self.optimizer_model.hook(self.model.parameters())
-        self.optimizer_weights.hook(self.weights_dict.parameters())
+        self.optimizer_weights.hook(self.weights.parameters())
+
+        # Add unknown parameters to optimization list in case of InverseProblem
         if isinstance(self.problem, InverseProblem):
             self.optimizer_model.instance.add_param_group(
                 {
@@ -216,110 +311,88 @@ def configure_optimizers(self):
                     ]
                 }
             )
+
+        # Hook the schedulers to the optimizers
         self.scheduler_model.hook(self.optimizer_model)
         self.scheduler_weights.hook(self.optimizer_weights)
+
         return (
             [self.optimizer_model.instance, self.optimizer_weights.instance],
             [self.scheduler_model.instance, self.scheduler_weights.instance],
         )
 
-    def on_train_start(self):
+    def _optimization_cycle(self, batch, batch_idx, **kwargs):
         """
-        This method is called at the start of the training process to set the
-        self-adaptive weights as parameters of the mask model.
-
-        :raises NotImplementedError: If the batch size is not ``None``.
-        """
-        if self.trainer.batch_size is not None:
-            raise NotImplementedError(
-                "SelfAdaptivePINN only works with full "
-                "batch size, set batch_size=None inside "
-                "the Trainer to use the solver."
-            )
-        device = torch.device(
-            self.trainer._accelerator_connector._accelerator_flag
-        )
+        Aggregate the loss for each condition in the batch.
 
-        # Initialize the self adaptive weights only for training points
-        for (
-            condition_name,
-            tensor,
-        ) in self.trainer.data_module.train_dataset.input.items():
-            self.weights_dict[condition_name].sa_weights.data = torch.rand(
-                (tensor.shape[0], 1), device=device
+        :param list[tuple[str, dict]] batch: A batch of data. Each element is a
+            tuple containing a condition name and a dictionary of points.
+        :param int batch_idx: The index of the current batch.
+        :param dict kwargs: Additional keyword arguments passed to
+            ``optimization_cycle``.
+        :return: The losses computed for all conditions in the batch, casted
+            to a subclass of :class:`torch.Tensor`. It should return a dict
+            containing the condition name and the associated scalar loss.
+        :rtype: dict
+        """
+        # Compute non-aggregated residuals
+        residuals = self.optimization_cycle(batch)
+
+        # Compute losses
+        losses = {}
+        for cond, res in residuals.items():
+
+            weight_tensor = self.weights[cond]()
+
+            # Get the correct indices for the weights. Modulus is used according
+            # to the number of points in the condition, as in the PinaDataset.
+            len_res = len(res)
+            idx = torch.arange(
+                batch_idx * len_res,
+                (batch_idx + 1) * len_res,
+                device=res.device,
+            ) % len(self.problem.input_pts[cond])
+
+            # Apply the weights to the residuals
+            losses[cond] = self._apply_reduction(
+                loss=(res * weight_tensor[idx])
             )
-        return super().on_train_start()
 
-    def on_load_checkpoint(self, checkpoint):
-        """
-        Override of the Pytorch Lightning ``on_load_checkpoint`` method to
-        handle checkpoints for Self-Adaptive Weights. This method should not be
-        overridden, if not intentionally.
-
-        :param dict checkpoint: Pytorch Lightning checkpoint dict.
-        """
-        # First initialize self-adaptive weights with correct shape,
-        # then load the values from the checkpoint.
-        for condition_name, _ in self.problem.input_pts.items():
-            shape = checkpoint["state_dict"][
-                f"_pina_models.1.{condition_name}.sa_weights"
-            ].shape
-            self.weights_dict[condition_name].sa_weights.data = torch.rand(
-                shape
+            # Store log
+            self.store_log(
+                f"{cond}_loss", losses[cond].item(), self.get_batch_size(batch)
             )
-        return super().on_load_checkpoint(checkpoint)
 
-    def loss_phys(self, samples, equation):
-        """
-        Computes the physics loss for the physics-informed solver based on the
-        provided samples and equation.
+        # Clamp unknown parameters in InverseProblem (if needed)
+        self._clamp_params()
 
-        :param LabelTensor samples: The samples to evaluate the physics loss.
-        :param EquationInterface equation: The governing equation.
-        :return: The computed physics loss.
-        :rtype: LabelTensor
-        """
-        residual = self.compute_residual(samples, equation)
-        weights = self.weights_dict[self.current_condition_name].forward()
-        loss_value = self._vectorial_loss(
-            torch.zeros_like(residual, requires_grad=True), residual
-        )
-        return self._vect_to_scalar(weights * loss_value)
-
-    def loss_data(self, input, target):
-        """
-        Compute the data loss for the PINN solver by evaluating the loss
-        between the network's output and the true solution. This method should
-        not be overridden, if not intentionally.
+        # Aggregate
+        loss = self.weighting.aggregate(losses).as_subclass(torch.Tensor)
 
-        :param input: The input to the neural network.
-        :type input: LabelTensor
-        :param target: The target to compare with the network's output.
-        :type target: LabelTensor
-        :return: The supervised loss, averaged over the number of observations.
-        :rtype: LabelTensor
-        """
-        return self._loss_fn(self.forward(input), target)
+        return loss
 
-    def _vect_to_scalar(self, loss_value):
+    def _apply_reduction(self, loss):
         """
-        Computation of the scalar loss.
+        Apply the specified reduction to the loss. The reduction is deferred
+        until the end of the optimization cycle to allow self-adaptive weights
+        to be applied to each point beforehand.
 
-        :param LabelTensor loss_value: the tensor of pointwise losses.
-        :raises RuntimeError: If the loss reduction is not ``mean`` or ``sum``.
-        :return: The computed scalar loss.
-        :rtype: LabelTensor
-        """
-        if self.loss.reduction == "mean":
-            ret = torch.mean(loss_value)
-        elif self.loss.reduction == "sum":
-            ret = torch.sum(loss_value)
-        else:
-            raise RuntimeError(
-                f"Invalid reduction, got {self.loss.reduction} "
-                "but expected mean or sum."
-            )
-        return ret
+        :param torch.Tensor loss: The loss tensor to be reduced.
+        :return: The reduced loss tensor.
+        :rtype: torch.Tensor
+        :raises ValueError: If the reduction method is neither "mean" nor "sum".
+        """
+        # Apply the specified reduction method
+        if self._reduction == "mean":
+            return loss.mean()
+        if self._reduction == "sum":
+            return loss.sum()
+
+        # Raise an error if the reduction method is not recognized
+        raise ValueError(
+            f"Unknown reduction: {self._reduction}."
+            " Supported reductions are 'mean' and 'sum'."
+        )
 
     @property
     def model(self):
@@ -332,7 +405,7 @@ def model(self):
         return self.models[0]
 
     @property
-    def weights_dict(self):
+    def weights(self):
         """
         The self-adaptive weights.
 
diff --git a/pina/utils.py b/pina/utils.py
index 569ba632c..ddbd2e8ac 100644
--- a/pina/utils.py
+++ b/pina/utils.py
@@ -47,15 +47,43 @@ class or, if ``subclass=True``, whether it is a subclass of the specified
         object_ = [object_]
 
     for obj in object_:
-        try:
-            if not subclass:
-                assert isinstance(obj, object_instance)
-            else:
-                assert issubclass(obj, object_instance)
-        except AssertionError as e:
-            raise ValueError(
-                f"{type(obj).__name__} must be {object_instance}."
-            ) from e
+        is_class = isinstance(obj, type)
+        expected_type_name = (
+            object_instance.__name__
+            if isinstance(object_instance, type)
+            else str(object_instance)
+        )
+
+        if subclass:
+            if not is_class:
+                raise ValueError(
+                    f"You passed {repr(obj)} "
+                    f"(an instance of {type(obj).__name__}), "
+                    f"but a {expected_type_name} class was expected. "
+                    f"Please pass a {expected_type_name} class or a "
+                    "derived one."
+                )
+            if not issubclass(obj, object_instance):
+                raise ValueError(
+                    f"You passed {obj.__name__} class, but a "
+                    f"{expected_type_name} class was expected. "
+                    f"Please pass a {expected_type_name} class or a "
+                    "derived one."
+                )
+        else:
+            if is_class:
+                raise ValueError(
+                    f"You passed {obj.__name__} class, but a "
+                    f"{expected_type_name} instance was expected. "
+                    f"Please pass a {expected_type_name} instance."
+                )
+            if not isinstance(obj, object_instance):
+                raise ValueError(
+                    f"You passed {repr(obj)} "
+                    f"(an instance of {type(obj).__name__}), "
+                    f"but a {expected_type_name} instance was expected. "
+                    f"Please pass a {expected_type_name} instance."
+                )
 
 
 def labelize_forward(forward, input_variables, output_variables):
diff --git a/tests/test_blocks/test_pirate_network_block.py b/tests/test_blocks/test_pirate_network_block.py
new file mode 100644
index 000000000..b827d24aa
--- /dev/null
+++ b/tests/test_blocks/test_pirate_network_block.py
@@ -0,0 +1,53 @@
+import torch
+import pytest
+from pina.model.block import PirateNetBlock
+
+data = torch.rand((20, 3))
+
+
+@pytest.mark.parametrize("inner_size", [10, 20])
+def test_constructor(inner_size):
+
+    PirateNetBlock(inner_size=inner_size, activation=torch.nn.Tanh)
+
+    # Should fail if inner_size is negative
+    with pytest.raises(AssertionError):
+        PirateNetBlock(inner_size=-1, activation=torch.nn.Tanh)
+
+
+@pytest.mark.parametrize("inner_size", [10, 20])
+def test_forward(inner_size):
+
+    model = PirateNetBlock(inner_size=inner_size, activation=torch.nn.Tanh)
+
+    # Create dummy embedding
+    dummy_embedding = torch.nn.Linear(data.shape[1], inner_size)
+    x = dummy_embedding(data)
+
+    # Create dummy U and V tensors
+    U = torch.rand((data.shape[0], inner_size))
+    V = torch.rand((data.shape[0], inner_size))
+
+    output_ = model(x, U, V)
+    assert output_.shape == (data.shape[0], inner_size)
+
+
+@pytest.mark.parametrize("inner_size", [10, 20])
+def test_backward(inner_size):
+
+    model = PirateNetBlock(inner_size=inner_size, activation=torch.nn.Tanh)
+    data.requires_grad_()
+
+    # Create dummy embedding
+    dummy_embedding = torch.nn.Linear(data.shape[1], inner_size)
+    x = dummy_embedding(data)
+
+    # Create dummy U and V tensors
+    U = torch.rand((data.shape[0], inner_size))
+    V = torch.rand((data.shape[0], inner_size))
+
+    output_ = model(x, U, V)
+
+    loss = torch.mean(output_)
+    loss.backward()
+    assert data.grad.shape == data.shape
diff --git a/tests/test_callback/test_optimizer_callback.py b/tests/test_callback/test_optimizer_callback.py
index 785a9c3f4..3383c792c 100644
--- a/tests/test_callback/test_optimizer_callback.py
+++ b/tests/test_callback/test_optimizer_callback.py
@@ -1,45 +1,63 @@
-from pina.callback import SwitchOptimizer
 import torch
 import pytest
 
 from pina.solver import PINN
 from pina.trainer import Trainer
 from pina.model import FeedForward
-from pina.problem.zoo import Poisson2DSquareProblem as Poisson
 from pina.optim import TorchOptimizer
+from pina.callback import SwitchOptimizer
+from pina.problem.zoo import Poisson2DSquareProblem as Poisson
+
+
+# Define the problem
+problem = Poisson()
+problem.discretise_domain(10)
+model = FeedForward(len(problem.input_variables), len(problem.output_variables))
+
+# Define the optimizer
+optimizer = TorchOptimizer(torch.optim.Adam)
 
-# make the problem
-poisson_problem = Poisson()
-boundaries = ["g1", "g2", "g3", "g4"]
-n = 10
-poisson_problem.discretise_domain(n, "grid", domains=boundaries)
-poisson_problem.discretise_domain(n, "grid", domains="D")
-model = FeedForward(
-    len(poisson_problem.input_variables), len(poisson_problem.output_variables)
-)
+# Initialize the solver
+solver = PINN(problem=problem, model=model, optimizer=optimizer)
 
-# make the solver
-solver = PINN(problem=poisson_problem, model=model)
+# Define new optimizers for testing
+lbfgs = TorchOptimizer(torch.optim.LBFGS, lr=1.0)
+adamW = TorchOptimizer(torch.optim.AdamW, lr=0.01)
 
-adam = TorchOptimizer(torch.optim.Adam, lr=0.01)
-lbfgs = TorchOptimizer(torch.optim.LBFGS, lr=0.001)
 
+@pytest.mark.parametrize("epoch_switch", [5, 10])
+@pytest.mark.parametrize("new_opt", [lbfgs, adamW])
+def test_switch_optimizer_constructor(new_opt, epoch_switch):
 
-def test_switch_optimizer_constructor():
-    SwitchOptimizer(adam, epoch_switch=10)
+    # Constructor
+    SwitchOptimizer(new_optimizers=new_opt, epoch_switch=epoch_switch)
 
+    # Should fail if epoch_switch is less than 1
+    with pytest.raises(ValueError):
+        SwitchOptimizer(new_optimizers=new_opt, epoch_switch=0)
 
-def test_switch_optimizer_routine():
-    # check initial optimizer
+
+@pytest.mark.parametrize("epoch_switch", [5, 10])
+@pytest.mark.parametrize("new_opt", [lbfgs, adamW])
+def test_switch_optimizer_routine(new_opt, epoch_switch):
+
+    # Check if the optimizer is initialized correctly
     solver.configure_optimizers()
-    assert solver.optimizer.instance.__class__ == torch.optim.Adam
-    # make the trainer
-    switch_opt_callback = SwitchOptimizer(lbfgs, epoch_switch=3)
+
+    # Initialize the trainer
+    switch_opt_callback = SwitchOptimizer(
+        new_optimizers=new_opt, epoch_switch=epoch_switch
+    )
     trainer = Trainer(
         solver=solver,
-        callbacks=[switch_opt_callback],
+        callbacks=switch_opt_callback,
         accelerator="cpu",
-        max_epochs=5,
+        max_epochs=epoch_switch + 2,
     )
     trainer.train()
-    assert solver.optimizer.instance.__class__ == torch.optim.LBFGS
+
+    # Check that the trainer strategy optimizers have been updated
+    assert solver.optimizer.instance.__class__ == new_opt.instance.__class__
+    assert (
+        trainer.strategy.optimizers[0].__class__ == new_opt.instance.__class__
+    )
diff --git a/tests/test_equations/test_system_equation.py b/tests/test_equations/test_system_equation.py
index 4a0a1163e..bf6268148 100644
--- a/tests/test_equations/test_system_equation.py
+++ b/tests/test_equations/test_system_equation.py
@@ -1,4 +1,4 @@
-from pina.equation import SystemEquation
+from pina.equation import SystemEquation, FixedValue, FixedGradient
 from pina.operator import grad, laplacian
 from pina import LabelTensor
 import torch
@@ -24,34 +24,78 @@ def foo():
     pass
 
 
-def test_constructor():
-    SystemEquation([eq1, eq2])
-    SystemEquation([eq1, eq2], reduction="sum")
+@pytest.mark.parametrize("reduction", [None, "mean", "sum"])
+def test_constructor(reduction):
+
+    # Constructor with callable functions
+    SystemEquation([eq1, eq2], reduction=reduction)
+
+    # Constructor with Equation instances
+    SystemEquation(
+        [
+            FixedValue(value=0.0, components=["u1"]),
+            FixedGradient(value=0.0, components=["u2"]),
+        ],
+        reduction=reduction,
+    )
+
+    # Constructor with mixed types
+    SystemEquation(
+        [
+            FixedValue(value=0.0, components=["u1"]),
+            eq1,
+        ],
+        reduction=reduction,
+    )
+
+    # Non-standard reduction not implemented
     with pytest.raises(NotImplementedError):
         SystemEquation([eq1, eq2], reduction="foo")
+
+    # Invalid input type
     with pytest.raises(ValueError):
         SystemEquation(foo)
 
 
-def test_residual():
+@pytest.mark.parametrize("reduction", [None, "mean", "sum"])
+def test_residual(reduction):
 
+    # Generate random points and output
     pts = LabelTensor(torch.rand(10, 2), labels=["x", "y"])
     pts.requires_grad = True
     u = torch.pow(pts, 2)
     u.labels = ["u1", "u2"]
 
-    eq_1 = SystemEquation([eq1, eq2], reduction="mean")
-    res = eq_1.residual(pts, u)
-    assert res.shape == torch.Size([10])
+    # System with callable functions
+    system_eq = SystemEquation([eq1, eq2], reduction=reduction)
+    res = system_eq.residual(pts, u)
+
+    # Checks on the shape of the residual
+    shape = torch.Size([10, 3]) if reduction is None else torch.Size([10])
+    assert res.shape == shape
 
-    eq_1 = SystemEquation([eq1, eq2], reduction="sum")
-    res = eq_1.residual(pts, u)
-    assert res.shape == torch.Size([10])
+    # System with Equation instances
+    system_eq = SystemEquation(
+        [
+            FixedValue(value=0.0, components=["u1"]),
+            FixedGradient(value=0.0, components=["u2"]),
+        ],
+        reduction=reduction,
+    )
 
-    eq_1 = SystemEquation([eq1, eq2], reduction=None)
-    res = eq_1.residual(pts, u)
-    assert res.shape == torch.Size([10, 3])
+    # Checks on the shape of the residual
+    shape = torch.Size([10, 3]) if reduction is None else torch.Size([10])
+    assert res.shape == shape
+
+    # System with mixed types
+    system_eq = SystemEquation(
+        [
+            FixedValue(value=0.0, components=["u1"]),
+            eq1,
+        ],
+        reduction=reduction,
+    )
 
-    eq_1 = SystemEquation([eq1, eq2])
-    res = eq_1.residual(pts, u)
-    assert res.shape == torch.Size([10, 3])
+    # Checks on the shape of the residual
+    shape = torch.Size([10, 3]) if reduction is None else torch.Size([10])
+    assert res.shape == shape
diff --git a/tests/test_model/test_pirate_network.py b/tests/test_model/test_pirate_network.py
new file mode 100644
index 000000000..f552f819d
--- /dev/null
+++ b/tests/test_model/test_pirate_network.py
@@ -0,0 +1,120 @@
+import torch
+import pytest
+from pina.model import PirateNet
+from pina.model.block import FourierFeatureEmbedding
+
+data = torch.rand((20, 3))
+
+
+@pytest.mark.parametrize("inner_size", [10, 20])
+@pytest.mark.parametrize("n_layers", [1, 3])
+@pytest.mark.parametrize("output_dimension", [2, 4])
+def test_constructor(inner_size, n_layers, output_dimension):
+
+    # Loop over the default and custom embedding
+    for embedding in [None, torch.nn.Linear(data.shape[1], inner_size)]:
+
+        # Constructor
+        model = PirateNet(
+            input_dimension=data.shape[1],
+            inner_size=inner_size,
+            output_dimension=output_dimension,
+            embedding=embedding,
+            n_layers=n_layers,
+            activation=torch.nn.Tanh,
+        )
+
+        # Check the default embedding
+        if embedding is None:
+            assert isinstance(model.embedding, FourierFeatureEmbedding)
+            assert model.embedding.sigma == 2.0
+
+        # Should fail if input_dimension is negative
+        with pytest.raises(AssertionError):
+            PirateNet(
+                input_dimension=-1,
+                inner_size=inner_size,
+                output_dimension=output_dimension,
+                embedding=embedding,
+                n_layers=n_layers,
+                activation=torch.nn.Tanh,
+            )
+
+        # Should fail if inner_size is negative
+        with pytest.raises(AssertionError):
+            PirateNet(
+                input_dimension=data.shape[1],
+                inner_size=-1,
+                output_dimension=output_dimension,
+                embedding=embedding,
+                n_layers=n_layers,
+                activation=torch.nn.Tanh,
+            )
+
+        # Should fail if output_dimension is negative
+        with pytest.raises(AssertionError):
+            PirateNet(
+                input_dimension=data.shape[1],
+                inner_size=inner_size,
+                output_dimension=-1,
+                embedding=embedding,
+                n_layers=n_layers,
+                activation=torch.nn.Tanh,
+            )
+
+        # Should fail if n_layers is negative
+        with pytest.raises(AssertionError):
+            PirateNet(
+                input_dimension=data.shape[1],
+                inner_size=inner_size,
+                output_dimension=output_dimension,
+                embedding=embedding,
+                n_layers=-1,
+                activation=torch.nn.Tanh,
+            )
+
+
+@pytest.mark.parametrize("inner_size", [10, 20])
+@pytest.mark.parametrize("n_layers", [1, 3])
+@pytest.mark.parametrize("output_dimension", [2, 4])
+def test_forward(inner_size, n_layers, output_dimension):
+
+    # Loop over the default and custom embedding
+    for embedding in [None, torch.nn.Linear(data.shape[1], inner_size)]:
+
+        model = PirateNet(
+            input_dimension=data.shape[1],
+            inner_size=inner_size,
+            output_dimension=output_dimension,
+            embedding=embedding,
+            n_layers=n_layers,
+            activation=torch.nn.Tanh,
+        )
+
+        output_ = model(data)
+        assert output_.shape == (data.shape[0], output_dimension)
+
+
+@pytest.mark.parametrize("inner_size", [10, 20])
+@pytest.mark.parametrize("n_layers", [1, 3])
+@pytest.mark.parametrize("output_dimension", [2, 4])
+def test_backward(inner_size, n_layers, output_dimension):
+
+    # Loop over the default and custom embedding
+    for embedding in [None, torch.nn.Linear(data.shape[1], inner_size)]:
+
+        model = PirateNet(
+            input_dimension=data.shape[1],
+            inner_size=inner_size,
+            output_dimension=output_dimension,
+            embedding=embedding,
+            n_layers=n_layers,
+            activation=torch.nn.Tanh,
+        )
+
+        data.requires_grad_()
+        output_ = model(data)
+
+        loss = torch.mean(output_)
+        loss.backward()
+        assert data.grad.shape == data.shape
diff --git a/tests/test_problem_zoo/test_inverse_poisson_2d_square.py b/tests/test_problem_zoo/test_inverse_poisson_2d_square.py
index 20a60e636..4304c8a5e 100644
--- a/tests/test_problem_zoo/test_inverse_poisson_2d_square.py
+++ b/tests/test_problem_zoo/test_inverse_poisson_2d_square.py
@@ -1,12 +1,25 @@
 from pina.problem.zoo import InversePoisson2DSquareProblem
 from pina.problem import InverseProblem, SpatialProblem
+import pytest
 
 
-def test_constructor():
-    problem = InversePoisson2DSquareProblem()
+@pytest.mark.parametrize("load", [True, False])
+@pytest.mark.parametrize("data_size", [0.01, 0.05])
+def test_constructor(load, data_size):
+
+    # Define the problem with or without loading data
+    problem = InversePoisson2DSquareProblem(load=load, data_size=data_size)
+
+    # Discretise the domain
     problem.discretise_domain(n=10, mode="random", domains="all")
+
+    # Check if the problem is correctly set up
     assert problem.are_all_domains_discretised
     assert isinstance(problem, InverseProblem)
     assert isinstance(problem, SpatialProblem)
     assert hasattr(problem, "conditions")
     assert isinstance(problem.conditions, dict)
+
+    # Should fail if data_size is not in the range [0.0, 1.0]
+    with pytest.raises(ValueError):
+        problem = InversePoisson2DSquareProblem(load=load, data_size=3.0)
diff --git a/tests/test_solver/test_competitive_pinn.py b/tests/test_solver/test_competitive_pinn.py
index 741390e31..8f585f029 100644
--- a/tests/test_solver/test_competitive_pinn.py
+++ b/tests/test_solver/test_competitive_pinn.py
@@ -20,14 +20,9 @@
 # define problems
 problem = Poisson()
 problem.discretise_domain(10)
-inverse_problem = InversePoisson()
+inverse_problem = InversePoisson(load=True, data_size=0.01)
 inverse_problem.discretise_domain(10)
 
-# reduce the number of data points to speed up testing
-data_condition = inverse_problem.conditions["data"]
-data_condition.input = data_condition.input[:10]
-data_condition.target = data_condition.target[:10]
-
 # add input-output condition to test supervised learning
 input_pts = torch.rand(10, len(problem.input_variables))
 input_pts = LabelTensor(input_pts, problem.input_variables)
diff --git a/tests/test_solver/test_gradient_pinn.py b/tests/test_solver/test_gradient_pinn.py
index 6e6c76c65..c28fc347e 100644
--- a/tests/test_solver/test_gradient_pinn.py
+++ b/tests/test_solver/test_gradient_pinn.py
@@ -31,14 +31,9 @@ class DummyTimeProblem(TimeDependentProblem):
 # define problems
 problem = Poisson()
 problem.discretise_domain(10)
-inverse_problem = InversePoisson()
+inverse_problem = InversePoisson(load=True, data_size=0.01)
 inverse_problem.discretise_domain(10)
 
-# reduce the number of data points to speed up testing
-data_condition = inverse_problem.conditions["data"]
-data_condition.input = data_condition.input[:10]
-data_condition.target = data_condition.target[:10]
-
 # add input-output condition to test supervised learning
 input_pts = torch.rand(10, len(problem.input_variables))
 input_pts = LabelTensor(input_pts, problem.input_variables)
diff --git a/tests/test_solver/test_pinn.py b/tests/test_solver/test_pinn.py
index ee501d876..d726047ef 100644
--- a/tests/test_solver/test_pinn.py
+++ b/tests/test_solver/test_pinn.py
@@ -20,14 +20,9 @@
 # define problems
 problem = Poisson()
 problem.discretise_domain(10)
-inverse_problem = InversePoisson()
+inverse_problem = InversePoisson(load=True, data_size=0.01)
 inverse_problem.discretise_domain(10)
 
-# reduce the number of data points to speed up testing
-data_condition = inverse_problem.conditions["data"]
-data_condition.input = data_condition.input[:10]
-data_condition.target = data_condition.target[:10]
-
 # add input-output condition to test supervised learning
 input_pts = torch.rand(10, len(problem.input_variables))
 input_pts = LabelTensor(input_pts, problem.input_variables)
diff --git a/tests/test_solver/test_rba_pinn.py b/tests/test_solver/test_rba_pinn.py
index 8eaf340ba..b464f3a7c 100644
--- a/tests/test_solver/test_rba_pinn.py
+++ b/tests/test_solver/test_rba_pinn.py
@@ -19,14 +19,9 @@
 # define problems
 problem = Poisson()
 problem.discretise_domain(10)
-inverse_problem = InversePoisson()
+inverse_problem = InversePoisson(load=True, data_size=0.01)
 inverse_problem.discretise_domain(10)
 
-# reduce the number of data points to speed up testing
-data_condition = inverse_problem.conditions["data"]
-data_condition.input = data_condition.input[:10]
-data_condition.target = data_condition.target[:10]
-
 # add input-output condition to test supervised learning
 input_pts = torch.rand(10, len(problem.input_variables))
 input_pts = LabelTensor(input_pts, problem.input_variables)
@@ -42,10 +37,14 @@
 @pytest.mark.parametrize("eta", [1, 0.001])
 @pytest.mark.parametrize("gamma", [0.5, 0.9])
 def test_constructor(problem, eta, gamma):
-    with pytest.raises(AssertionError):
-        solver = RBAPINN(model=model, problem=problem, gamma=1.5)
     solver = RBAPINN(model=model, problem=problem, eta=eta, gamma=gamma)
 
+    with pytest.raises(ValueError):
+        solver = RBAPINN(model=model, problem=problem, gamma=1.5)
+
+    with pytest.raises(ValueError):
+        solver = RBAPINN(model=model, problem=problem, eta=-0.1)
+
     assert solver.accepted_conditions_types == (
         InputTargetCondition,
         InputEquationCondition,
@@ -54,30 +53,18 @@ def test_constructor(problem, eta, gamma):
 
 
 @pytest.mark.parametrize("problem", [problem, inverse_problem])
-def test_wrong_batch(problem):
-    with pytest.raises(NotImplementedError):
-        solver = RBAPINN(model=model, problem=problem)
-        trainer = Trainer(
-            solver=solver,
-            max_epochs=2,
-            accelerator="cpu",
-            batch_size=10,
-            train_size=1.0,
-            val_size=0.0,
-            test_size=0.0,
-        )
-        trainer.train()
-
-
-@pytest.mark.parametrize("problem", [problem, inverse_problem])
+@pytest.mark.parametrize("batch_size", [None, 1, 5, 20])
 @pytest.mark.parametrize("compile", [True, False])
-def test_solver_train(problem, compile):
-    solver = RBAPINN(model=model, problem=problem)
+@pytest.mark.parametrize(
+    "loss", [torch.nn.L1Loss(reduction="sum"), torch.nn.MSELoss()]
+)
+def test_solver_train(problem, batch_size, loss, compile):
+    solver = RBAPINN(model=model, problem=problem, loss=loss)
     trainer = Trainer(
         solver=solver,
         max_epochs=2,
         accelerator="cpu",
-        batch_size=None,
+        batch_size=batch_size,
         train_size=1.0,
         val_size=0.0,
         test_size=0.0,
@@ -89,14 +76,18 @@ def test_solver_train(problem, compile):
 
 
 @pytest.mark.parametrize("problem", [problem, inverse_problem])
+@pytest.mark.parametrize("batch_size", [None, 1, 5, 20])
 @pytest.mark.parametrize("compile", [True, False])
-def test_solver_validation(problem, compile):
-    solver = RBAPINN(model=model, problem=problem)
+@pytest.mark.parametrize(
+    "loss", [torch.nn.L1Loss(reduction="sum"), torch.nn.MSELoss()]
+)
+def test_solver_validation(problem, batch_size, loss, compile):
+    solver = RBAPINN(model=model, problem=problem, loss=loss)
     trainer = Trainer(
         solver=solver,
         max_epochs=2,
         accelerator="cpu",
-        batch_size=None,
+        batch_size=batch_size,
         train_size=0.9,
         val_size=0.1,
         test_size=0.0,
@@ -108,14 +99,18 @@ def test_solver_validation(problem, compile):
 
 
 @pytest.mark.parametrize("problem", [problem, inverse_problem])
+@pytest.mark.parametrize("batch_size", [None, 1, 5, 20])
 @pytest.mark.parametrize("compile", [True, False])
-def test_solver_test(problem, compile):
-    solver = RBAPINN(model=model, problem=problem)
+@pytest.mark.parametrize(
+    "loss", [torch.nn.L1Loss(reduction="sum"), torch.nn.MSELoss()]
+)
+def test_solver_test(problem, batch_size, loss, compile):
+    solver = RBAPINN(model=model, problem=problem, loss=loss)
     trainer = Trainer(
         solver=solver,
         max_epochs=2,
         accelerator="cpu",
-        batch_size=None,
+        batch_size=batch_size,
         train_size=0.7,
         val_size=0.2,
         test_size=0.1,
diff --git a/tests/test_solver/test_self_adaptive_pinn.py b/tests/test_solver/test_self_adaptive_pinn.py
index aba43da5f..b2d1361ca 100644
--- a/tests/test_solver/test_self_adaptive_pinn.py
+++ b/tests/test_solver/test_self_adaptive_pinn.py
@@ -20,14 +20,9 @@
 # define problems
 problem = Poisson()
 problem.discretise_domain(10)
-inverse_problem = InversePoisson()
+inverse_problem = InversePoisson(load=True, data_size=0.01)
 inverse_problem.discretise_domain(10)
 
-# reduce the number of data points to speed up testing
-data_condition = inverse_problem.conditions["data"]
-data_condition.input = data_condition.input[:10]
-data_condition.target = data_condition.target[:10]
-
 # add input-output condition to test supervised learning
 input_pts = torch.rand(10, len(problem.input_variables))
 input_pts = LabelTensor(input_pts, problem.input_variables)
@@ -42,9 +37,11 @@
 @pytest.mark.parametrize("problem", [problem, inverse_problem])
 @pytest.mark.parametrize("weight_fn", [torch.nn.Sigmoid(), torch.nn.Tanh()])
 def test_constructor(problem, weight_fn):
+
+    solver = SAPINN(problem=problem, model=model, weight_function=weight_fn)
+
     with pytest.raises(ValueError):
         SAPINN(model=model, problem=problem, weight_function=1)
-    solver = SAPINN(problem=problem, model=model, weight_function=weight_fn)
 
     assert solver.accepted_conditions_types == (
         InputTargetCondition,
@@ -53,26 +50,13 @@ def test_constructor(problem, weight_fn):
     )
 
 
-@pytest.mark.parametrize("problem", [problem, inverse_problem])
-def test_wrong_batch(problem):
-    with pytest.raises(NotImplementedError):
-        solver = SAPINN(model=model, problem=problem)
-        trainer = Trainer(
-            solver=solver,
-            max_epochs=2,
-            accelerator="cpu",
-            batch_size=10,
-            train_size=1.0,
-            val_size=0.0,
-            test_size=0.0,
-        )
-        trainer.train()
-
-
 @pytest.mark.parametrize("problem", [problem, inverse_problem])
 @pytest.mark.parametrize("compile", [True, False])
-def test_solver_train(problem, compile):
-    solver = SAPINN(model=model, problem=problem)
+@pytest.mark.parametrize(
+    "loss", [torch.nn.L1Loss(reduction="sum"), torch.nn.MSELoss()]
+)
+def test_solver_train(problem, compile, loss):
+    solver = SAPINN(model=model, problem=problem, loss=loss)
     trainer = Trainer(
         solver=solver,
         max_epochs=2,
@@ -95,8 +79,11 @@ def test_solver_train(problem, compile):
 
 @pytest.mark.parametrize("problem", [problem, inverse_problem])
 @pytest.mark.parametrize("compile", [True, False])
-def test_solver_validation(problem, compile):
-    solver = SAPINN(model=model, problem=problem)
+@pytest.mark.parametrize(
+    "loss", [torch.nn.L1Loss(reduction="sum"), torch.nn.MSELoss()]
+)
+def test_solver_validation(problem, compile, loss):
+    solver = SAPINN(model=model, problem=problem, loss=loss)
     trainer = Trainer(
         solver=solver,
         max_epochs=2,
@@ -119,8 +106,11 @@ def test_solver_validation(problem, compile):
 
 @pytest.mark.parametrize("problem", [problem, inverse_problem])
 @pytest.mark.parametrize("compile", [True, False])
-def test_solver_test(problem, compile):
-    solver = SAPINN(model=model, problem=problem)
+@pytest.mark.parametrize(
+    "loss", [torch.nn.L1Loss(reduction="sum"), torch.nn.MSELoss()]
+)
+def test_solver_test(problem, compile, loss):
+    solver = SAPINN(model=model, problem=problem, loss=loss)
     trainer = Trainer(
         solver=solver,
         max_epochs=2,