Make xCEBRA compatible with the batched inference & padding in solver

Author: CeliaBenquet

cebra/solver/base.py

@@ -599,7 +599,7 @@ def transform(self,
             session_id: The session ID, an :py:class:`int` between 0 and
                 the number of sessions -1 for multisession, and set to
                 ``None`` for single session.
-            batch_size: If not None, batched inference will be applied.
+            batch_size: If not None, batched inference will not be applied.
 
         Returns:
             The output embedding.

cebra/solver/multiobjective.py

@@ -53,6 +53,7 @@
 import cebra.data
 import cebra.io
 import cebra.models
+import cebra.solver.single_session as cebra_solver_single
 from cebra.solver import register
 from cebra.solver.base import Solver
 from cebra.solver.schedulers import Scheduler
@@ -187,7 +188,7 @@ def _process_info(self, info):
 
 
 @dataclasses.dataclass
-class MultiobjectiveSolverBase(Solver):
+class MultiobjectiveSolverBase(cebra_solver_single.SingleSessionSolver):
 
     feature_ranges: List[slice] = None
     renormalize: bool = None
@@ -209,6 +210,13 @@ def __post_init__(self):
             renormalize=self.renormalize,
         )
 
+    def parameters(self, session_id: Optional[int] = None):
+        """Iterate over all parameters."""
+        super().parameters(session_id=session_id)
+
+        for parameter in self.regularizer.parameters():
+            yield parameter
+
     def fit(self,
             loader: cebra.data.Loader,
             valid_loader: cebra.data.Loader = None,
@@ -241,6 +249,7 @@ def _run_validation():
                 save_hook(solver=self, step=num_steps)
             return stats_val
 
+        self._set_fitted_params(loader)
         self.to(loader.device)
 
         iterator = self._get_loader(loader,
@@ -393,11 +402,14 @@ def validation(
         logger=None,
         weights_loss: Optional[List[float]] = None,
     ):
+        loader.dataset.configure_for(self.model)
+        iterator = self._get_loader(loader)
+
         self.model.eval()
         total_loss = Meter()
 
         losses_dict = {}
-        for _, batch in enumerate(loader):
+        for _, batch in iterator:
             predictions = self._inference(batch)
             losses = self.criterion(predictions)
 
@@ -445,7 +457,7 @@ def validation(
         return stats_val
 
     @torch.no_grad()
-    def transform(self, inputs: torch.Tensor) -> torch.Tensor:
+    def transform_deprecated(self, inputs: torch.Tensor) -> torch.Tensor:
         offset = self.model.get_offset()
         self.model.eval()
         X = inputs.cpu().numpy()