How to save the tensor in the forward pass of customer autograd function


Is there have any method to save a tensor for statistics?


class Exp(Function):
     def forward(ctx, i):
         i_max = i.max()
         return i

    def backward(ctx, grad_output):
        return grad_output



If you want to save it for the backward pass, then you can save it within the context as ctx.i_max = i.max().
If you want to save it for external use, then you can use any pythonic way you want: use a global variable, save it in a list, have a stats object where you can record what happens etc.

1 Like