flambe.learn.utils

Module Contents

flambe.learn.utils.select_device(device: Optional[str]) → str[source]

Chooses the torch device to run in.

Parameters
device: Union[torch.device, str]
A device or a string representing a device, such as ‘cpu’
str
the passed-as-parameter device if any, otherwise cuda if available. Last option is cpu.