torch.cuda.memory.memory_allocated

torch.cuda.memory.memory_allocated(device=None)

返回给定设备上张量当前占用的 GPU 内存（以字节为单位）。

参数

device (torch.device or int, optional) – selected device. Returns statistic for the current device, given by current_device(), if device is None (default).

返回类型

int

注意

This is likely less than the amount shown in nvidia-smi since some unused memory can be held by the caching allocator and some context needs to be created on GPU. See Memory management for more details about GPU memory management.