Skip to content

Instantly share code, notes, and snippets.

@Stonesjtu
Stonesjtu / mem_report.py
Last active March 7, 2023 16:58
A simple Pytorch memory usages profiler
import gc
import torch
## MEM utils ##
def mem_report():
'''Report the memory usage of the tensor.storage in pytorch
Both on CPUs and GPUs are reported'''
def _mem_report(tensors, mem_type):