starling.utilities.check_device

check_device(use_device, default_device='gpu')

Function to check the device was correctly set.

Parameters:

• use_device (str) – Identifier for the device to be used for predictions. Possible inputs: ‘cpu’, ‘mps’, ‘cuda’, ‘cuda:int’, where the int corresponds to the index of a specific cuda-enabled GPU. If ‘cuda’ is specified and cuda.is_available() returns False, this will raise an Exception If ‘mps’ is specified and mps is not available, an exception will be raised.

• default_device (str) – The default device to use if device=None. If device=None and default_device != ‘cpu’ and default_device is not available, device_string will be returned as ‘cpu’. Default is ‘gpu’. This checks first for cuda and then for mps because STARLING is faster on both than it is on CPU, so we should use the fastest device available. Options are ‘cpu’ or ‘gpu’

Returns:

torch.device

Return type:

A PyTorch device object representing the device to use.