torch.cuda.memory_allocated

torch.cuda.memory_allocated(device=None)

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

參數

device (torch.device 或 int, optional) – 選擇的裝置。如果 device 為 None (預設值)，則傳回由 current_device() 給定的當前裝置的統計資訊。

回傳類型

int

注意

這可能小於 nvidia-smi 中顯示的數量，因為快取分配器可能會保留一些未使用的記憶體，並且需要在 GPU 上建立一些上下文。有關 GPU 記憶體管理的更多詳細資訊，請參閱 記憶體管理。