Print info about allocated tensors in outer scopes

torch_tensor_info.py

import gc
import inspect
import torch

def _print_tensors(depth=2):
    """Print info about tensor variables locally and nonlocally."""
    print('='*80)
    for obj in gc.get_objects():
        frame = inspect.currentframe()
        frames = []
        for i in range(depth):
            frame = frame.f_back
            frames.append(frame)
        try:
            if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)):
                for gobj, gobjid in globals().items():
                    if gobjid is obj and gobj not in ('obj', 'gobjid'):
                        print('GLOBAL:\t', gobj, obj.size(), obj.data_ptr())
                for frame in frames:
                    fname = frame.f_code.co_name
                    locs = frame.f_locals
                    for gobj, gobjid in locs.items():
                        if gobjid is obj and gobj not in ('obj', 'gobjid'):
                            print(f'{fname}():\t', gobj, obj.size(), obj.data_ptr())
        except:
            pass
    print('='*80)