Does not work with torch 1.7.1+

[alexus37]

The demo code does not work with this version of PyTorch. No output is printed

import torch
from pytorch_memlab import LineProfiler

def inner():
    torch.nn.Linear(100, 100).cuda()

def outer():
    linear = torch.nn.Linear(100, 100).cuda()
    linear2 = torch.nn.Linear(100, 100).cuda()
    inner()

with LineProfiler(outer, inner) as prof:
    outer()
    prof.display()

[alexus37]

sorry, I was too fast, it works but the documentation is a bit confusing. the last line should be prof.print_stats()