set_device

dragon.cuda.set_device(device_index=0)[source]

Set the current device.

Parameters:
  • device_index (int, optional, default=0) – The device index.