Add learning rate scheduler

[hmosousa]

Cosine annealing is the way to go: https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.CosineAnnealingLR.html