Device

metatensor.models.utils.devices.pick_devices(architecture_devices: List[str], desired_device: str | None = None) List[device][source]

Pick (best) devices for training.

The choice is made on the intersection of the architecture_devices and the available devices on the current system. If no desired_device is provided the first device of this intersection will be returned.

Parameters:
  • architecture_devices (List[str]) – Devices supported by the architecture. The list should be sorted by the preference of the architecture while the most prefferred device should be first and the least one last.

  • desired_device (str | None) – desired device by the user

Return type:

List[device]