Fix cosine schedule
Fixing similar problem to #327 and #324 in cosine schedule.

Btw, do you think it would make sense to have [something like this](https://github.com/lukovnikov/pytorch-pretrained-BERT/blob/optim/pytorch_pretrained_bert/optimization.py) for both of your `optimization.py`'s?
Hi @lukovnikov, yes I think the various schedules you have added to your fork are very nice!
Do you want to add them in this PR as well?
Otherwise, I'll merge it.<|||||>Merging it for now. Thanks @lukovnikov <|||||>Hi, sorry, lost track of this, will make a new PR soon.