device

dragon.device(
  device_type,
  device_index=0
)[source]

Context-manager to nest the the device spec.

Examples:

with dragon.device('cuda', 0):
    x = dragon.constant(1)
Parameters:
  • device_type ({'cpu', 'gpu', 'cuda', 'cnml'}, required) – The type of device.
  • device_index (int, optional, default=0) – The index of the device.
Returns:

Dict – The current default device spec.