LION.utils.device

The device LION uses when none is given.

Pick the GPU once, at the top of a script, and LION follows it:

torch.cuda.set_device(2) # or torch.set_default_device(“cuda:2”)

Anything given an explicit device= uses that instead. On machines with many GPUs where only one may be used, the most robust option is to make only that GPU visible to the process: CUDA_VISIBLE_DEVICES=2 python script.py.

Functions

default_device()

Return the device LION uses when no device is given.


LION.utils.device.default_device()[source]

Return the device LION uses when no device is given.

In order of priority:

  1. torch.set_default_device(...), if set to anything other than the CPU;

  2. the current CUDA device (torch.cuda.set_device(k), cuda:0 by default), if CUDA is available;

  3. the CPU.