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"