arogozhnikov:

Unfortunately,  I don’t see where this requirement (store only input or output) comes from and how do I get around it (storing as a member in Function means you can use this function only once).
Any ideas how to store intermediate things?


I see two main options.

Using the autograd.Function instance only once is a great way to do this. In fact this is what used to happen all the time when you used operations on variables and is the right thing to do™.
If you absolutely dislike autograd.Function instances or - like myself - like to feel modern, go for the new-style autograd.Function and change the line with the function application to  y_pred = MyReLU.apply(x.mm(w1)).mm(w2). You do not need relu when doing this. This will store inputs and other stuff in the context ctx and is what happens nowadays when you use operations on variables.

Best regards
Thomas
class MyReLU(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return input.clamp(min=0)

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_variables
        grad_input = grad_output.clone()
        grad_input[input < 0] = 0
        return grad_input