Pytorch Print Tensors

print_tensors.py

import torch
import gc
        
def tensor_meta_data(tensor):
    element_count = 1;
    for dim in tensor.size():
        element_count = element_count * dim
    size_in_bytes = element_count * tensor.element_size()
    dtype = str(tensor.dtype).replace("torch.", "")
    size = str(tensor.size()).replace("torch.Size(", "").replace(")", "")
    return f"{size_in_bytes/1000000:5.1f}MB {dtype}{size} {type(tensor).__name__} {tensor.device}"

def print_tensors():
    for obj in gc.get_objects():
        if torch.is_tensor(obj):
            print(tensor_meta_data(obj))

Helper function to print the size and memory consumption of tensors in pytorch to help debug out-of-memory errors.