hot fix for MultiStepLR_HotFix

Author: ildoonet

FastAutoAugment/lr_scheduler.py

@@ -1,5 +1,5 @@
 import torch
-
+from torch.optim.lr_scheduler import MultiStepLR
 from theconf import Config as C
 
 
@@ -10,8 +10,14 @@ def adjust_learning_rate_resnet(optimizer):
     """
 
     if C.get()['epoch'] == 90:
-        return torch.optim.lr_scheduler.MultiStepLR(optimizer, [30, 60, 80])
+        return MultiStepLR_HotFix(optimizer, [30, 60, 80])
     elif C.get()['epoch'] == 270:   # autoaugment
-        return torch.optim.lr_scheduler.MultiStepLR(optimizer, [90, 180, 240])
+        return MultiStepLR_HotFix(optimizer, [90, 180, 240])
     else:
         raise ValueError('invalid epoch=%d for resnet scheduler' % C.get()['epoch'])
+
+        
+class MultiStepLR_HotFix(MultiStepLR):
+    def __init__(self, optimizer, milestones, gamma=0.1, last_epoch=-1):
+        super(MultiStepLR_HotFix, self).__init__(optimizer, milestones, gamma, last_epoch)
+        self.milestones = list(milestones)