From 18bdaa8f6f054367920676767be5c23219489be4 Mon Sep 17 00:00:00 2001 From: jiltseb Date: Fri, 24 May 2024 19:39:23 +0000 Subject: [PATCH] added docstring --- faster_whisper/transcribe.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/faster_whisper/transcribe.py b/faster_whisper/transcribe.py index 468669f5..b768c3a2 100644 --- a/faster_whisper/transcribe.py +++ b/faster_whisper/transcribe.py @@ -184,11 +184,16 @@ def _sanitize_parameters(self, **kwargs): preprocess_kwargs["maybe_arg"] = kwargs["maybe_arg"] return preprocess_kwargs, {}, {} - def get_device(self, device): + def get_device(self, device: Union[int, str, "torch.device"]): """ - Returns a torch.device object given a device of the form Union[int, str, "torch.device"]. - The function also makes vad_device same as self.model.device by default ("auto"). - + Converts the input device into a torch.device object. + + The input can be an integer, a string, or a `torch.device` object. + + The function handles a special case where the input device is "auto". + When "auto" is specified, the device will default to the + device of the model (self.model.device). If the model's device is also "auto", + it selects "cuda" if a CUDA-capable device is available; otherwise, it selects "cpu". """ if isinstance(device, torch.device): return device