I architected my codebase so that modules don’t know the data types of other modules it may reference. The only structs shared between modules are my own internal structs.
So I have a struct called ManagedMemory which handles memory management. So what I usually do is instantiate a ManagedMemory(dataSize, dataPtr, memoryType=cuda/cpu) inside the ModulePytorch from the inference output and then send that through pubsub to another module. ManagedMemory is now the struct which owns the data.
Just a thought which may work (tell me if this is doable):

Maybe an alternative would be to up the refCount artificially by 1 and then pass in a custom function into my ManagedMemory which will decrease the refCount when it’s not being used anymore.
ManagedMemory(dataSize, dataPtr, memoryType=cuda/cpu, [&]() { // Decrease Tensor Refcount })
Then that means the data is still managed by Module PyTorch but will have a floating pointer until the refCount falls to 0.

I just don’t like the idea of doing this. I’m more of favor of telling tensors to release ownership to the pointers.