How to identify tensors in autograd saved tensors hook?

1

I learned about Hooks for autograd saved tensors — PyTorch Tutorials 2.6.0+cu124 documentation.

But I am missing one detail. I’d like to implement a pack/unpack hook which applies to specific nodes only. I’d expect that when the tensor is created in forward, I can tag it somehow and then identify it in pack_hook. But I failed to find anything.

What is the simplest way to identify specific tensors in the pack hook?

2

My first thought was to use grad_fn.name to identify the tensors, also to tag some extra information that I need in the pack_hook. But I can only access this once the node is created, and this means pack_hook will be called already, right?

3

I guess I found a solution to this, but it is pretty indirect. I push the annotation into a list and check the list in the hook, popping the annotation if there is one.