@dfalbel
I think this depends on where you call the malloc, either in main thread or other threads.
‘https://github.com/pytorch/pytorch/blob/master/c10/cuda/CUDACachingAllocator.cpp#L214 5’
The memory free callback is triggered whenever you call malloc().