import torch
import torch.nn as nn
from torch.optim.lr_scheduler import ConstantLR , ReduceLROnPlateau, SequentialLR


model = nn.Sequential(nn.Linear(2,50),
                      nn.Sigmoid(),
                      nn.Linear(50,50),
                      nn.Sigmoid(),
                      nn.Linear(50,2)
                      )


criteria = nn.CrossEntropyLoss()
optim =  torch.optim.SGD(model.parameters(),lr=.1,momentum=.9)

def get_toyData(N1,N2):
    labels = torch.cat([torch.ones(N1),torch.zeros(N2)])
    data = torch.cat([torch.randn(N1,2),torch.randn(N2,2)+1.5])
    return data , labels 

N1 , N2 = 1000,1000 
train_data , train_labels = get_toyData(N1,N2)

dataset = torch.utils.data.TensorDataset(train_data,train_labels)
dl = torch.utils.data.DataLoader(dataset,batch_size=10,
                                 shuffle=True)

scheduler1 = ConstantLR(optim, factor=0.1, total_iters=40)
#scheduler2 = ExponentialLR(optim, gamma=0.9)
scheduler2 = ReduceLROnPlateau(optim, 'min')
    
for i in range(100):
    totalLoss = 0
    train_error = 0
    for x,t in dl:
        optim.zero_grad()
        out=model(x)
        loss = criteria(out,t.long())
        loss.backward()
        optim.step()
        totalLoss += loss.item()
    if i <40:
        
        scheduler1.step()
    else:
        scheduler2.step(totalLoss)
    print(i , optim.state_dict()['param_groups'][0]['lr'])
        
    #print(i,totalLoss,train_error/1300)


Assuming optimizer uses lr = 0.1 for all groups
lr = 0.01     if epoch < 40
lr = 0.1    if epoch =40
and ReduceLROnPlateau by .1  if epoch > 40