torch.cuda.set_device¶
- torch.cuda.set_device(device)[source]¶
Sets the current device.
Usage of this function is discouraged in favor of
device
. In most cases it’s better to useCUDA_VISIBLE_DEVICES
environmental variable.- Parameters:
device (torch.device or int) – selected device. This function is a no-op if this argument is negative.