Segmentation fault when using torch.profiler

autograd
1

Hi, everyone. I come across following problem using torch.profiler.

from torch.profiler import profile, ProfilerActivity

    with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
        a = torch.rand(1).cuda()
        a += 1
        print('***Profiling...***) # This line will not be reached if call cuda() on Tensor a.
        train_model(model, datasets, cfg)
    print(prof.key_averages().table(sort_by='cuda_memory_usage'))

When I run the code above, segmentation fault is raised without any other error information.
If cuda() is not called on Tensor a, the next print function works.

The code works fine without profiling.
I’ll appreciate if anyone point the reasons out.