qimpy.set_gpu_visibility

set_gpu_visibility(local_rank)

Update CUDA_VISIBLE_DEVICES to select one GPU based on local_rank of process. Return the device number of the selected GPU, and -1 if no GPUs specified. (Note that CUDA_VISIBLE_DEVICES must be set explicitly to use GPUs.)

Parameters:

local_rank (int) –

Return type:

int