from l4 import L4

#...

optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)                                                           | create mode 100644 l4.py                                                                                                                  
# wrap original optimizer with L4
l4opt = L4(optimizer) 

#...

loss = F.nll_loss(output, target)
loss.backward() 

# Comment out original optimizer step
# optimizer.step()

# make step with L4 optimizer, dont forget to pass loss value
l4opt.step(loss)