Fix import errors

Author: CeliaBenquet

cebra/data/multi_session.py

@@ -26,12 +26,12 @@
 
 import literate_dataclasses as dataclasses
 import torch
+import torch.nn as nn
 
 import cebra.data as cebra_data
 import cebra.distributions
 from cebra.data.datatypes import Batch
 from cebra.data.datatypes import BatchIndex
-from cebra.models import Model
 
 __all__ = [
     "MultiSessionDataset",
@@ -105,7 +105,7 @@ def load_batch(self, index: BatchIndex) -> List[Batch]:
             ) for session_id, session in enumerate(self.iter_sessions())
         ]
 
-    def configure_for(self, model: "Model"):
+    def configure_for(self, model: "cebra.models.Model"):
         """Configure the dataset offset for the provided model.
 
         Call this function before indexing the dataset. This sets the
@@ -114,9 +114,16 @@ def configure_for(self, model: "Model"):
         Args:
             model: The model to configure the dataset for.
         """
-        self.offset = model.get_offset()
-        for session in self.iter_sessions():
-            session.configure_for(model)
+        if not isinstance(model, nn.ModuleList):
+            raise ValueError(
+                "The model must be a nn.ModuleList to configure the dataset.")
+        if len(model) != self.num_sessions:
+            raise ValueError(
+                f"The model must have {self.num_sessions} sessions, but got {len(model)}."
+            )
+
+        for i, session in enumerate(self.iter_sessions()):
+            session.configure_for(model[i])
 
 
 @dataclasses.dataclass

tests/_utils_deprecated.py

@@ -9,7 +9,6 @@
 import cebra
 import cebra.integrations.sklearn.utils as sklearn_utils
 import cebra.models
-import cebra.solvers
 
 
 #NOTE: Deprecated: transform is now handled in the solver but the original
@@ -79,7 +78,7 @@ def cebra_transform_deprecated(cebra_model,
 #       using the transform method of the model, and handling padding is implemented
 #       directly in the base Solver. This method is kept for testing purposes.
 @torch.no_grad()
-def multiobjective_transform_deprecated(solver: cebra.solvers.Solver,
+def multiobjective_transform_deprecated(solver: "cebra.solvers.Solver",
                                         inputs: torch.Tensor) -> torch.Tensor:
     """Transform the input data using the model.