fix: replace use_gpu with accelerator and devices (#24)

BREAKING CHANGE: users should now pass in `accelerator` and `devices` to
`train` instead of `use_gpu`, in line with the changes introduced in
scvi-tools 1.1.

velovi/_model.py

@@ -14,7 +14,7 @@
 from scvi.dataloaders import DataSplitter
 from scvi.model.base import BaseModelClass, UnsupervisedTrainingMixin, VAEMixin
 from scvi.train import TrainingPlan, TrainRunner
-from scvi.utils._docstrings import setup_anndata_dsp
+from scvi.utils._docstrings import devices_dsp, setup_anndata_dsp
 
 from ._constants import REGISTRY_KEYS
 from ._module import VELOVAE
@@ -118,12 +118,14 @@ def __init__(
         )
         self.init_params_ = self._get_init_params(locals())
 
+    @devices_dsp.dedent
     def train(
         self,
         max_epochs: Optional[int] = 500,
         lr: float = 1e-2,
         weight_decay: float = 1e-2,
-        use_gpu: Optional[Union[str, int, bool]] = None,
+        accelerator: str = "auto",
+        devices: Union[int, list[int], str] = "auto",
         train_size: float = 0.9,
         validation_size: Optional[float] = None,
         batch_size: int = 256,
@@ -143,9 +145,8 @@ def train(
             Learning rate for optimization
         weight_decay
             Weight decay for optimization
-        use_gpu
-            Use default GPU if available (if None or True), or index of GPU to use (if int),
-            or name of GPU (if str, e.g., `'cuda:0'`), or use CPU (if False).
+        %(param_accelerator)s
+        %(param_devices)s
         train_size
             Size of training set in the range [0.0, 1.0].
         validation_size
@@ -189,7 +190,8 @@ def train(
             training_plan=training_plan,
             data_splitter=data_splitter,
             max_epochs=max_epochs,
-            use_gpu=use_gpu,
+            accelerator=accelerator,
+            devices=devices,
             **trainer_kwargs,
         )
         return runner()