Source code for flambe.learn.utils
from typing import Optional
import torch
[docs]def select_device(device: Optional[str]) -> str:
"""
Chooses the torch device to run in.
Parameters
------------
device: Union[torch.device, str]
A device or a string representing a device, such as 'cpu'
Returns
---------
str
the passed-as-parameter device if any, otherwise
cuda if available. Last option is cpu.
"""
if device is not None:
return device
else:
return "cuda" if torch.cuda.is_available() else "cpu"