Appreciated for your reply!
m_o = ModelOutputs(net, layer=40)
output, conv_output = m_o(input)
loss = cross_entropy(output, label)
loss.backward()
grad_m_o = m_o.gradients
grad_net = net.features[40].weight.grad

I find the grad_m_o  is different from grad_net, why?
And I got  another proplem: I want to load model in DataParallel and use multi-GPU in one server. like this:
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2'
device = 'cuda' if torch.cuda.is_available() else 'cpu'

... ...
net.to(device)
net= torch.nn.DataParallel(model)
input.to(device) 
...

And I use the first approach in the code above.

class ModelOutputs():

m_o = ModelOutputs(net, layer)
output, conv_output = m_o(input)

I passed the net in ModelOutpus. When I ran the code, it did’t use multi-GPU but gpu’0.
Is the problem caused by not defined as nn.Module?