Fix offset in docs

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:`cebra_data.Dataset.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:`cebra_data.Dataset.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

@@ -91,14 +91,15 @@ def _check_indices(batch_start_idx: int, batch_end_idx: int,
 
 
 def _add_batched_zero_padding(batched_data: torch.Tensor,
-                              offset: cebra.data.Offset, batch_start_idx: int,
+                              offset: cebra.data.Offset, 
+                              batch_start_idx: int,
                               batch_end_idx: int,
                               num_samples: int) -> torch.Tensor:
     """Add zero padding to the input data before inference.
 
     Args:
         batched_data: Data to apply the inference on.
-        offset (cebra.data.Offset): _description_
+        offset: Offset of the model to consider when padding.
         batch_start_idx: Index of the first sample in the batch.
         batch_end_idx: Index of the first sample in the batch.
         num_samples (int): Total number of samples in the data.