From f6bbbabd8e6cba95ef3892f74814d4722da99a0f Mon Sep 17 00:00:00 2001 From: Philipp Weiler Date: Sun, 17 Mar 2024 11:20:26 +0100 Subject: [PATCH] Update `VELOVI::train` Remove argument `use_gpu` and rely on `accelerator` and `devices`, instead. --- scvelo/tools/_vi_model.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/scvelo/tools/_vi_model.py b/scvelo/tools/_vi_model.py index 856cd8be..c2f56223 100644 --- a/scvelo/tools/_vi_model.py +++ b/scvelo/tools/_vi_model.py @@ -127,7 +127,6 @@ def train( 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, @@ -149,9 +148,14 @@ 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). + accelerator + Supports passing different accelerator types `("cpu", "gpu", "tpu", "ipu", "hpu", "mps, "auto")` as well as + custom accelerator instances. + devices + The devices to use. Can be set to a non-negative index (`int` or `str`), a sequence of device indices + (`list` or comma-separated `str`), the value `-1` to indicate all available devices, or `"auto"` for + automatic selection based on the chosen `accelerator`. If set to `"auto"` and `accelerator` is not + determined to be `"cpu"`, then `devices` will be set to the first available device. train_size Size of training set in the range [0.0, 1.0]. validation_size @@ -195,7 +199,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()