torch.accelerator.memory.get_memory_info

torch.accelerator.memory.get_memory_info(device_index=None, /)

返回指定设备索引的当前设备内存信息。

参数:

device_index (torch.device, str, int, optional) – 要定位的设备索引。如果未给出，则默认使用 torch.accelerator.current_device_index()。如果提供了 torch.device 或 str，则其类型必须与当前的 accelerator 设备类型匹配。

返回:

返回一个包含两个整数的元组 (free_memory, total_memory)，单位为字节。

第一个值是设备上的空闲内存（适用于所有进程和应用程序），第二个值是设备的硬件总内存容量。

返回类型:

tuple[int, int]