torch.cuda.memory.memory_allocated

torch.cuda.memory.memory_allocated(device=None) [исходный код]

Возвращает текущий объем памяти GPU в байтах, занятый тензорами на указанном устройстве.

Параметры:

device (torch.device или int, необязательно) – выбранное устройство. Возвращает статистику для текущего устройства, определяемого с помощью current_device(), если device имеет значение None (по умолчанию).

Тип возвращаемого значения:

int

Примечание

Это значение, вероятно, меньше, чем объем памяти, показанный в nvidia-smi, поскольку часть неиспользуемой памяти может удерживаться кэширующим аллокатором, а на GPU необходимо создать некоторый контекст. Подробнее об управлении памятью GPU см. в разделе Управление памятью.