Fix wrong learning rate evaluation in CosineAnnealingLR in Python 2 (#4656)

nguyen-binh-minh · soumith · commit d9001ce8619e · 2018-02-07T08:58:33.000-08:00
diff --git a/test/test_optim.py b/test/test_optim.py
@@ -465,7 +465,7 @@ def test_cos_anneal_lr(self):
         epochs = 10
         eta_min = 1e-10
         single_targets = [eta_min + (0.05 - eta_min) *
-                          (1 + math.cos(x / epochs * math.pi)) / 2
+                          (1 + math.cos(math.pi * x / epochs)) / 2
                           for x in range(epochs)]
         targets = [single_targets, list(map(lambda x: x * epochs, single_targets))]
         scheduler = CosineAnnealingLR(self.opt, T_max=epochs, eta_min=eta_min)
diff --git a/torch/optim/lr_scheduler.py b/torch/optim/lr_scheduler.py
@@ -194,7 +194,7 @@ def __init__(self, optimizer, T_max, eta_min=0, last_epoch=-1):
 
     def get_lr(self):
         return [self.eta_min + (base_lr - self.eta_min) *
-                (1 + math.cos(self.last_epoch / self.T_max * math.pi)) / 2
+                (1 + math.cos(math.pi * self.last_epoch / self.T_max)) / 2
                 for base_lr in self.base_lrs]
 
 
