I am having issues with my learning rate not adjusting properly. Instead of steadily decreasing, it is going from the initial learning rate to 0 repeatedly.

This is the code for my scheduler:

lrs = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max = len(train_loader))

and this is how it is being implemented in my training loop:
optimizer.zero_grad()
  loss.backward()
  optimizer.step()
  lrs.step()

How can I fix this?