samlab.torch module

Functionality for simplifying work with Pytorch.

samlab.torch.select_device()

Automatically choose a Torch device to be used for subsequent computation.

Automatically chooses GPUs over CPUs, prioritizing GPUs that are lightly loaded.

Note

Always returns a valid device, but there is no guarantee that a given device will have sufficient resources for a given computation.