Update VELOVI::train

Remove `use_gpu` from DataSplitter.
theislab · Jun 8, 2023 · a7b7cf7 · a7b7cf7
1 parent 8aaa334
commit a7b7cf7
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/scvelo/tools/_vi_model.py b/scvelo/tools/_vi_model.py
@@ -126,6 +126,8 @@ def train(
         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,
@@ -179,7 +181,6 @@ def train(
             train_size=train_size,
             validation_size=validation_size,
             batch_size=batch_size,
-            use_gpu=use_gpu,
         )
         training_plan = TrainingPlan(self.module, **plan_kwargs)