# Layer weights in custom Loss results in autograd problem

Hello,

I want to introduce a kind of total variation penalty.

I want to realize this by the fact, that the neural network is penalized for it, if the neighboured weights differ strongly.

But I want to apply this only in one layer. So I have added this additionally in the loss function.

My code parts are therefore as follows:

``````class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()

# decoder
self.conv1 = nn.Conv1d(in_channels=1, out_channels=1, kernel_size=50, bias=False)
self.conv2 = nn.BatchNorm1d(1)
self.conv3 = nn.Conv1d(in_channels=1,out_channels=1, kernel_size=30, bias=False)
self.conv4 = nn.BatchNorm1d(1)
self.conv5 = nn.Conv1d(in_channels=1,out_channels=1, kernel_size=11, bias=False)
self.conv6 = nn.BatchNorm1d(1)
self.ASC = nn.Softmax(dim=1)
# encoder
self.DEC = nn.Linear(in_features=3,out_features=91)

def forward(self, x):
x = F.leaky_relu(self.conv1(x))
x = F.leaky_relu(self.conv2(x))
x = F.leaky_relu(self.conv3(x))
x = F.leaky_relu(self.conv4(x))
x = F.leaky_relu(self.conv5(x))
x = x*3.5
x = F.leaky_relu(self.conv6(x))
a = self.ASC(x)
x = F.leaky_relu(self.DEC(a))
return x, a
net = CNN()
``````
``````def train(net, trainloader, NUM_EPOCHS, cur_epoch=0):
train_loss = []
for epoch in range(cur_epoch, NUM_EPOCHS):
running_loss = 0.0
inputData = data
inputData = inputData.to(device)
outputData ,a = net(inputData)
loss_MSE = criterion(outputData, inputData)
loss_TVPenalty = TV_Penalty(net.DEC)
loss = loss_MSE + loss_TVPenalty
loss.backward()
optimizer.step()
net.DEC.weight.data.clamp_(min=0)
running_loss += loss.item()
train_loss.append(loss)
return train_loss
``````
``````def TV_Penalty(Layer):
loss = torch.pow((Layer.weight.data[1:]-Layer.weight.data[:-1]), 2).sum()*TV_PENALTY_CONSTANT
return loss

``````

If I use it this way, the autograde function of Pytorch is not applied to the `python def TV_Penalty(Layer): ` function.

How to use the autograd in Pytocht on the structure of the weights of the network itself? So then this is considered in an additional loss function.

You should be able to use the `.weight` parameters directly, instead of their `.data` attribute, which Autograd cannot track and which should thus be avoided.

Hihih, many thanks, works like a charm.