get_device#

composer.utils.get_device(device)[source]#

Takes string or Device and returns the corresponding Device.

Parameters: device (str | Device, optional) – A string corresponding to a device (one of 'cpu', 'gpu', 'mps', or 'tpu') or a Device.
Returns: Device – Device corresponding to the passed string or Device. If no argument is passed, returns DeviceGPU if available, or DeviceCPU if no GPU is available.