Improve docs

theislab · Nov 11, 2024 · aab870e · aab870e
1 parent 5454956
commit aab870e
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 2 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -68,6 +68,7 @@ dev = [
     "twine>=4.0.2",
 ]
 doc = [
+    "merlin-dataloader==23.8.0",
     "docutils>=0.8,!=0.18.*,!=0.19.*",
     "sphinx>=4",
     "sphinx-book-theme>=1.0.0",

diff --git a/src/drvi/model/__init__.py b/src/drvi/model/__init__.py
@@ -1,3 +1,4 @@
 from drvi.scvi_tools_based.model import DRVI
+from drvi.scvi_tools_based.module import DRVIModule
 
-__all__ = ["DRVI"]
+__all__ = ["DRVI", "DRVIModule"]
diff --git a/src/drvi/scvi_tools_based/model/_drvi.py b/src/drvi/scvi_tools_based/model/_drvi.py
@@ -49,7 +49,7 @@ class DRVI(VAEMixin, DRVIArchesMixin, UnsupervisedTrainingMixin, BaseModelClass,
     categorical_covariates
         Categorical Covariates as a list of texts. You can specify emb dimension by appending @dim to each cpvariate.
     **model_kwargs
-        Keyword args for :class:`~drvi.model.DRVI`
+        Keyword args for :class:`~drvi.model.DRVIModule`
 
     Examples
     --------