add buffer and checks on eta

Author: GiovanniCanali

pina/solver/physics_informed_solver/rba_pinn.py

@@ -94,6 +94,7 @@ def __init__(
             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,
@@ -114,24 +115,20 @@ def __init__(
                 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 the weight of each point to 0
-        self.weights = {
-            cond: torch.zeros((len(data), 1), device=self.device)
-            for cond, data in self.problem.input_pts.items()
-        }
-
-    def on_train_start(self):
-        """
-        Hook method called at the beginning of training.
-        """
-        device = self.trainer.strategy.root_device
-        for cond in self.weights:
-            self.weights[cond] = self.weights[cond].to(device)
-        return super().on_train_start()
+        self.weights = {}
+        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}")
 
     def training_step(self, batch, batch_idx, **kwargs):
         """

tests/test_solver/test_rba_pinn.py

@@ -47,6 +47,9 @@ def test_constructor(problem, eta, 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,