From b9a5cfa8135ca624d07e0cb65f8e644c259a2a93 Mon Sep 17 00:00:00 2001
From: Martin Kim <46072231+martinkim0@users.noreply.github.com>
Date: Mon, 11 Mar 2024 10:39:33 -0700
Subject: [PATCH] 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 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/velovi/_model.py b/velovi/_model.py
index e75b98b..df8ee5d 100644
--- a/velovi/_model.py
+++ b/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()