add callbacks to docs

scverse · Feb 19, 2022 · c82aa51 · c82aa51
1 parent fb99341
commit c82aa51
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/docs/api/developer.rst b/docs/api/developer.rst
@@ -187,6 +187,8 @@ TrainingPlans define train/test/val optimization steps for modules.
    train.Trainer
    train.TrainingPlan
    train.TrainRunner
+   train.SaveBestState
+   train.LoudEarlyStopping
 
 
 Utilities

diff --git a/scvi/train/__init__.py b/scvi/train/__init__.py
@@ -1,4 +1,4 @@
-from ._callbacks import LoudEarlyStopping
+from ._callbacks import LoudEarlyStopping, SaveBestState
 from ._trainer import Trainer
 from ._trainingplans import (
     AdversarialTrainingPlan,
@@ -18,4 +18,5 @@
     "ClassifierTrainingPlan",
     "TrainRunner",
     "LoudEarlyStopping",
+    "SaveBestState",
 ]