src.common.get_device

src.common.get_device(device_str='auto')

Get PyTorch device

Args:

device_str: Device specification (“auto”, “cpu”, “cuda”, “cuda:0”, etc.)

Returns:

PyTorch device

Parameters:

device_str (str)

Return type:

torch.device