torch.accelerator.get_device_capability

torch.accelerator.get_device_capability(device=None, /)

Return the capability of the currently selected device.

Parameters:

device (torch.device, str, int, optional) – The device to query capabilities for accelerator device type. If not given, use torch.accelerator.current_device_index() by default.

Returns:

A dictionary containing device capability information. The dictionary includes:

• supported_dtypes (set(torch.dtype)): Set of PyTorch data types for which tensors can be allocated on the accelerator and type conversion across supported dtypes are supported. Any operator support outside of that is not guaranteed

Return type:

dict[str, Any]

Examples

>>> # Query capabilities for current device
>>> capabilities = torch.accelerator.get_device_capability("cuda:0")
>>> print("Supported dtypes:", capabilities["supported_dtypes"])