heat.devices

handle different devices. Current options: CPU (default), GPU

Module Contents

class Device(device_type: str, device_id: int, torch_device: str)

Implements a compute device. HeAT can run computations on different compute devices or backends. A device describes the device type and id on which said computation should be carried out.

Parameters:

• device_type (str) – Represents HeAT’s device name

• device_id (int) – The device id

• torch_device (str) – The corresponding PyTorch device type

Examples

>>> ht.Device("cpu", 0, "cpu:0")
device(cpu:0)
>>> ht.Device("gpu", 0, "cuda:0")
device(gpu:0)

__repr__() → str

Return the unambiguous information of Device.

__str__() → str

Return the descriptive information of Device.

__eq__(other: Any) → bool

Overloads the == operator for local equal check.

Parameters:

other (Any) – The object to compare with

cpu

The standard CPU Device

Examples

>>> ht.cpu
device(cpu:0)
>>> ht.ones((2, 3), device=ht.cpu)
DNDarray([[1., 1., 1.],
          [1., 1., 1.]], dtype=ht.float32, device=cpu:0, split=None)

get_device() → Device

Retrieves the currently globally set default Device.

sanitize_device(device: str | Device | None = None) → Device

Sanitizes a device or device identifier, i.e. checks whether it is already an instance of Device or a string with known device identifier and maps it to a proper Device.

Parameters:

device (str or Device, optional) – The device to be sanitized

Raises:

ValueError – If the given device id is not recognized

use_device(device: str | Device | None = None) → None

Sets the globally used default Device.

Parameters:

device (str or Device) – The device to be set