@dfalbel
I never did it before. But I think it is doable.
The FreeCudaMemoryCallbacksRegistry is defined here and used to register FreeMemoryCallbacks
‘https://github.com/pytorch/pytorch/blob/master/c10/cuda/CUDACachingAllocator.cpp#L23 3’
FreeCudaMemoryCallbacksRegistry is actually a c10::registry that defined here:
‘https://github.com/pytorch/pytorch/blob/master/c10/util/Registry.h#L268 2’
‘https://github.com/pytorch/pytorch/blob/master/c10/util/Registry.h#L215 2’
‘https://github.com/pytorch/pytorch/blob/master/c10/util/Registry.h#L268 2’
FreeMemoryCallback is defined here:
‘https://github.com/pytorch/pytorch/blob/6f396e18c33efbb8547b7ca9d412a316bd51429b/c10/cuda/CUDACachingAllocator.h#L20 3’
So to add your own free memory callback you can do the following:

Write your own FreeMemoryCallback by inheriting FreeMemoryCallback.
Implement your own execute()
Use FreeCudaMemoryCallbacksRegistry()->register(…) to register your own
see c10::Registry to learn what parameter should be passed in, it basically the
type of your customized memory callback class and a creator func ptr.

Your customized free memory callback will be called in trigger_free_memory_callbacks then.