fix: correctly initialize LRSchedulers for multi-gpu training
Previously the LRScheduler would be passed in the original optimizer,
and if horovod is being user, we would replace that optimzer with the
horovod optimizer, without updating the optimizer that the LRSCheduler
was using. Now we initialize the LRSCheduler with the correct optimizer.
#Done.
@@ -114,8 +114,6 @@ class PyTorchTrialController(det.LoopTrialController):
self.optimizer = self.trial.optimizer(self.model)
# TODO: Check that optimizer is not an amp optimizer.
- self.lr_helper = _LRHelper(self.trial.create_lr_scheduler(self.optimizer))
-
self._init_device()
self.model = self.model.to(self.device)
@@ -139,6 +137,8 @@ class PyTorchTrialController(det.LoopTrialController):
self.model = nn.DataParallel(self.model)
logging.debug("Initialized mode for native parallel training.")
+ self.lr_helper = _LRHelper(self.trial.create_lr_scheduler(self.optimizer))
+
# If a load path is provided load weights and restore the data location.
self._load()
