set_device¶

dragon.vm.torch.cuda.set_device(device)[source]¶

Set the current device.

Parameters:

• device (Union[dragon.vm.torch.device, int]) – The device to set.