From 2f564f65648ab6b31fba6828738f5e6abd3dacb4 Mon Sep 17 00:00:00 2001 From: Steffen Schneider Date: Sun, 20 Oct 2024 19:16:22 +0200 Subject: [PATCH] Fix docstrings --- cebra/integrations/matplotlib.py | 4 ++-- cebra/integrations/sklearn/cebra.py | 3 +-- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/cebra/integrations/matplotlib.py b/cebra/integrations/matplotlib.py index 1d4dee5..9bc1888 100644 --- a/cebra/integrations/matplotlib.py +++ b/cebra/integrations/matplotlib.py @@ -290,13 +290,13 @@ def _define_plot_dim( * If ``idx_order`` is not provided, the plot will be 3D by default. * If ``idx_order`` is provided, if it has 3 dimensions, the plot will be 3D, if only 2 dimensions - are provided, the plot will be 2D. + are provided, the plot will be 2D. If the embedding dimension is equal to 2: * If ``idx_order`` is not provided, the plot will be 2D by default. * If ``idx_order`` is provided, if it has 3 dimensions, the plot will be 3D, if 2 dimensions - are provided, the plot will be 2D. + are provided, the plot will be 2D. This is supposing that the dimensions provided to ``idx_order`` are in the range of the number of dimensions of the embedding (i.e., between 0 and :py:attr:`cebra.CEBRA.output_dimension` -1). diff --git a/cebra/integrations/sklearn/cebra.py b/cebra/integrations/sklearn/cebra.py index abacd53..046d334 100644 --- a/cebra/integrations/sklearn/cebra.py +++ b/cebra/integrations/sklearn/cebra.py @@ -1337,13 +1337,12 @@ def save(self, - 'state': The state of the CEBRA model, which includes various internal attributes. - 'state_dict': The state dictionary of the underlying solver used by CEBRA. - 'metadata': Additional metadata about the saved model, including the backend used and - the version of CEBRA PyTorch, NumPy and scikit-learn. + the version of CEBRA PyTorch, NumPy and scikit-learn. "torch" backend: The model is directly saved using `torch.save` with no additional information. The saved file contains the entire CEBRA model state. - Example: >>> import cebra