torch.get_device_module

torch.get_device_module(device=None)

Returns the module associated with a given device(e.g., torch.device(‘cuda’), “mtia:0”, “xpu”, …). If no device is given, return the module for the current accelerator or CPU if none is present.