# Add a layer derivative in the loss function

I am writing a NN in pytorch and I want to add the derivative of the output with respect to one of the inner layers in the loss. Here is a simple example of what I mean:

class NN(torch.nn.Module):

def __init__(self):
super(Model, self).__init__()
self.l1 = torch.nn.Linear(2, 3)
self.l2 = torch.nn.Linear(3, 1)
self.l3 = torch.nn.Linear(1, 3)
self.l4 = torch.nn.Linear(3, 2)

def forward(self, x):
z = torch.sigmoid(self.l1(x))
z = torch.sigmoid(self.l2(z))
y_pred = torch.sigmoid(self.l3(z))
y_pred = self.l4(y_pred)
return y_pred


In this case (for just one example) I want the loss to be defined as: loss = (y-y_pred)**2+norm(d(y_pred)/dz) (the last part is not correct from a python point of view). So what I want is to add the derivative of of the output with respect to a given layer to the loss function. I tried some things with hook, but that is not really what I need here. Can someone help me? Thank you!