Skip to content

Commit

Permalink
added docstring
Browse files Browse the repository at this point in the history
  • Loading branch information
Jiltseb committed May 24, 2024
1 parent 883be4d commit 18bdaa8
Showing 1 changed file with 9 additions and 4 deletions.
13 changes: 9 additions & 4 deletions faster_whisper/transcribe.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 18bdaa8

Please sign in to comment.