AttributeError in pytorch

Question

I want use https://github.com/mmbejani/TikhonovRegularizationTerm. This library includes the implementation of regularization tikhonov terms that were published from 2010 until now (2020).

I am trying to train a very basic feed forward nn with this lib on mnist data set and getting the following error:

AttributeError: 'WeightDecay' object has no attribute 'backward'

'

optimizer = optim.SGD(model.parameters(), lr=0.003, momentum=0.9)
time0 = time()
epochs = 15
for e in range(epochs):
    running_loss = 0
    for images, labels in trainloader:
        # Flatten MNIST images into a 784 long vector
        images = images.view(images.shape[0], -1)

        # Training pass
        optimizer.zero_grad()

        output = model(images)
        loss_function = nn.CrossEntropyLoss()
        loss_function_with_regularization = WeightDecay(model, loss_function)

        #This is where the model learns by backpropagating
        loss_function_with_regularization.backward()

        #And optimizes its weights here
        optimizer.step()

        running_loss += loss.item()
    else:
        print("Epoch {} - Training loss: {}".format(e,running_loss/len(trainloader)))
print("\nTraining Time (in minutes) =",(time()-time0)/60)