fix docs build, missing refs

Author: stes

cebra/data/datasets.py

@@ -398,7 +398,7 @@ def __getitem__(self, index):
         return self.neural[index].transpose(2, 1)
 
     def load_batch_supervised(self, index: Batch,
-                              labels_supervised) -> torch.tensor:
+                              labels_supervised) -> torch.Tensor:
         """Load a batch for supervised learning.
 
         Args:

cebra/data/single_session.py

@@ -172,10 +172,9 @@ class ContinuousDataLoader(cebra_data.Loader):
     * auxiliary variables, using the empirical distribution of how behavior various across
       ``time_offset`` timesteps (``time_delta``). Sampling for this setting is implemented
       in :py:class:`cebra.distributions.continuous.TimedeltaDistribution`.
-    * alternatively, the distribution can be selected to be a Gaussian or von Mises-Fisher distribution
+    * alternatively, the distribution can be selected to be a Gaussian distribution
       parametrized by a fixed ``delta`` around the reference sample, using the implementation in
-      :py:class:`cebra.distributions.continuous.DeltaNormalDistribution` and
-      :py:class:`cebra.distributions.continuous.DeltaVMFDistribution`.
+      :py:class:`cebra.distributions.continuous.DeltaNormalDistribution`.
 
     Args:
         See dataclass fields.
@@ -228,11 +227,13 @@ def _init_distribution(self):
                     self.dataset.continuous_index,
                     self.delta,
                     device=self.device)
-            elif self.conditional == "delta_vmf":
-                self.distribution = cebra.distributions.DeltaVMFDistribution(
-                    self.dataset.continuous_index,
-                    self.delta,
-                    device=self.device)
+            # TODO(stes): Add this distribution from internal xCEBRA codebase at a later point
+            # in time, currently not in use.
+            #elif self.conditional == "delta_vmf":
+            #    self.distribution = cebra.distributions.DeltaVMFDistribution(
+            #        self.dataset.continuous_index,
+            #        self.delta,
+            #        device=self.device)
             else:
                 raise ValueError(self.conditional)