torch.xpu.memory_allocated

torch.xpu.memory_allocated(device=None)

返回指定裝置上 tensors 目前佔用的 GPU 記憶體大小 (以位元組為單位)。

參數

device (torch.device 或 int 或 str, 可選) – 選定的裝置。如果 device 為 None (預設值)，則返回由 current_device() 提供的目前裝置的統計資訊。

返回類型

int

注意

此數值可能小於 xpu-smi 中顯示的數量，因為快取分配器可能會保留一些未使用的記憶體，並且需要在 GPU 上建立一些上下文。