branches, optim cosine fix

pytorch_pretrained_bert/optimization.py

@@ -26,7 +26,9 @@ logger = logging.getLogger(__name__)
 def warmup_cosine(x, warmup=0.002):
     if x < warmup:
         return x/warmup
-    return 0.5 * (1.0 + torch.cos(math.pi * x))
+
+    x_ = (x - warmup) / (1 - warmup)  # progress after warmup
+    return 0.5 * (1. + math.cos(math.pi * x_))
 
 def warmup_constant(x, warmup=0.002):
     """ Linearly increases learning rate over `warmup`*`t_total` (as provided to BertAdam) training steps.