torchcompat 1.1.4

torch compatibility layer

pip install torchcompat

Features

• Provide a super set implementation of pytorch device interface to enable code to run seamlessly between different accelerators.

• Identify uniquely devices

import torchcompat.core as accelerator

# on  cuda accelerator == torch.cuda
# on  rocm accelerator == torch.cuda
# on   xpu accelerator == torch.xpu
# on gaudi accelerator == ...

assert accelerator.is_available() == true
assert accelerator.device_name in ('xpu', 'cuda', "hpu")           # rocm is seen as cuda by pytorch
assert accelerator.device_string(0) == "cuda:0" or "xpu:0" or "hpu:0"
assert accelerator.fetch_device(0) == torch.device("cuda:0")


accelerator.set_enable_tf32(true) # toggle the right flags for each backend

Example

example here