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

Hi,

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

Example

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

    @staticmethod
    def backward(ctx, grad_output):
        return grad_output

Thanks!

Hi,

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