device¶

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

Context-manager to nest the device.

Examples:

with dragon.device('cuda', 0):
    x = dragon.constant(1)

Parameters:

• device_type (str) – The device type.
• device_index (int, optional, default=0) – The device index.

Returns:

dragon.DeviceSpec – The current nesting device spec.