Fix extra docs errors

Author: CeliaBenquet

cebra/data/multi_session.py

@@ -110,7 +110,7 @@ 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
-        :py:attr:`offset` attribute of the dataset.
+        :py:attr:`cebra_data.Dataset.offset` attribute of the dataset.
 
         Args:
             model: The model to configure the dataset for.

cebra/data/single_session.py

@@ -76,7 +76,7 @@ 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
-        :py:attr:`offset` attribute of the dataset.
+        :py:attr:`cebra_data.Dataset.offset` attribute of the dataset.
 
         Args:
             model: The model to configure the dataset for.

cebra/solver/base.py

@@ -185,7 +185,7 @@ def _transform(
     model: cebra.models.Model,
     inputs: torch.Tensor,
     pad_before_transform: bool,
-    offset: cebra.data.Offset,
+    offset: cebra.data.datatypes.Offset,
 ) -> torch.Tensor:
     """Compute the embedding.
 
@@ -206,7 +206,7 @@ def _transform(
 
 def _batched_transform(model: cebra.models.Model, inputs: torch.Tensor,
                        batch_size: int, pad_before_transform: bool,
-                       offset: cebra.data.Offset) -> torch.Tensor:
+                       offset: cebra.data.datatypes.Offset) -> torch.Tensor:
     """Compute the embedding on batched inputs.
 
     Args:

tests/test_solver.py

@@ -65,6 +65,7 @@
 
 # multi_session_tests.append((*args, cebra.solver.MultiSessionAuxVariableSolver))
 
+
 def _get_loader(data, loader_initfunc):
     kwargs = dict(num_steps=5, batch_size=32)
     loader = loader_initfunc(data, **kwargs)
@@ -574,6 +575,7 @@ def test_select_model_multi_session(data_name, model_name, session_id,
             assert offset.left == offset_.left and offset.right == offset_.right
             assert model == model_
 
+
 models = [
     "offset1-model",
     "offset10-model",
@@ -683,7 +685,7 @@ def test_batched_transform_multi_session(data_name, model_name, padding,
     n_samples = dataset._datasets[0].neural.shape[0]
     assert all(
         d.neural.shape[0] == n_samples for d in dataset._datasets
-    ), # all sessions need to have same number of samples
+    ), "for this set all of the sessions need to have same number of samples."
 
     smallest_batch_length = n_samples - batch_size
     offset_ = model[0].get_offset()