from torch.optim.lr_scheduler import LambdaLR
import torch.nn as nn
import torch
from torch.optim import *


class net(nn.Module):
    def __init__(self):
        super(net,self).__init__()
        self.fc = nn.Linear(1,10)
    def forward(self,x):
        return self.fc(x)
model = net()

optimizer = Adam(model.parameters(),lr = 0.01)
max_update_num=10
lambda1 = lambda epoch: (1 - epoch/max_update_num)
optimizer.scheduler = lr_scheduler.LambdaLR(optimizer,lr_lambda = lambda1)