torch.accelerator.current_accelerator — PyTorch 2.7 documentation (original) (raw)
torch.accelerator.current_accelerator(check_available=False)[source][source]¶
Return the device of the accelerator available at compilation time. If no accelerator were available at compilation time, returns None. See accelerator for details.
Parameters
check_available (bool, optional) – if True, will also do a runtime check to see if the device torch.accelerator.is_available() on top of the compile-time check. Default: False
Returns
return the current accelerator as torch.device.
Return type
Example:
If an accelerator is available, sent the model to it
model = torch.nn.Linear(2, 2) if (current_device := current_accelerator(check_available=True)) is not None: model.to(current_device)