Save the gradient of the intermediate tensor

I have a question, if I want to save the gradient of the intermediate tensor generated in backward for future use, what should I do? Below is a small demo. Thanks!

import torch
import gc
a = []
class Exp(torch.autograd.Function):
    @staticmethod
    def forward(ctx, i):
        result = torch.exp(i)
        ctx.save_for_backward(result)
        return result
    @staticmethod
    def backward(ctx, grad_output):
        result, = ctx.saved_tensors
        a.append(grad_output.detach())
        return grad_output * result
exp = Exp()
x = torch.tensor([3., 4.], requires_grad=True)
y = exp.apply(x)
y.sum().backward()
gc.collect()
torch.cuda.empty_cache()
print(a)

cc @ptrblck
A similar question to Use custom stream in custom torch.autograd.Function

Your custom Function idea sounds fine to me. You can also register a backward hook to y here, e.g. y.register_hook(fn) and do the same appending idea.

Alternatively, if you were okay until the end of backward to use this gradient, you can just specify it as part of the inputs= you pass to .grad()/.backward(), e.g. `y.sum().backward(inputs=(y, […])), this would mean however that you’d need to explicitly pass the rest of things you want to compute grad wrt to into inputs=.