henry000 commited on
Commit
ea0ad19
Β·
1 Parent(s): cdd9a11

πŸ› [Fix] unstable learing rate schedule

Browse files
Files changed (1) hide show
  1. yolo/tools/solver.py +1 -1
yolo/tools/solver.py CHANGED
@@ -83,7 +83,7 @@ class TrainModel(ValidateModel):
83
  return self.train_loader
84
 
85
  def on_train_epoch_start(self):
86
- self.trainer.optimizers[0].next_epoch(len(self.train_loader))
87
 
88
  def training_step(self, batch, batch_idx):
89
  lr_dict = self.trainer.optimizers[0].next_batch()
 
83
  return self.train_loader
84
 
85
  def on_train_epoch_start(self):
86
+ self.trainer.optimizers[0].next_epoch((len(self.train_loader) + 1) // self.trainer.world_size)
87
 
88
  def training_step(self, batch, batch_idx):
89
  lr_dict = self.trainer.optimizers[0].next_batch()