Hi jytug: If in addition, I want to pass some other information backward layer by layer besides the output_grad. what should I do ? There is no need to add anything on forward. I try to use the following code while doesn’t work. Thanks a lot !
import torch
import torch.autograd as autograd
import torch.nn as nn

class MyFun(torch.autograd.Function):
    def forward(self, inp):
        return inp

def backward(self, grad_out, P):
    grad_input = grad_out.clone()
    print('Custom backward called!')
    return grad_input, P-1

class MyMod(nn.Module):
    def forward(self, x):
        return MyFun()(x)

mod = nn.Sequential(MyMod(), MyMod())

y = autograd.Variable(torch.randn(1), requires_grad=True)
z = mod(y)
P = torch.ones((1,1))
z.backward(P)