utils.torch_utils module

Utilities for setting up PyTorch

get_device(cuda_ID=0, needs_cuda=False)[source]

Choose

Parameters

  • cuda_ID (int) – if CUDA is available, use device with this ID

  • needs_cuda (bool) – whether to raise an error if CUDA is not available

Returns

Return type

torch.device