henry000 commited on
Commit
8553686
Β·
1 Parent(s): 17274a5

πŸ› [Fix] the bug in lr scheduler, use ceil

Browse files
Files changed (1) hide show
  1. yolo/tools/solver.py +2 -1
yolo/tools/solver.py CHANGED
@@ -1,3 +1,4 @@
 
1
  from pathlib import Path
2
 
3
  from lightning import LightningModule
@@ -83,7 +84,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()
 
1
+ from math import ceil
2
  from pathlib import Path
3
 
4
  from lightning import LightningModule
 
84
  return self.train_loader
85
 
86
  def on_train_epoch_start(self):
87
+ self.trainer.optimizers[0].next_epoch(ceil(len(self.train_loader) / self.trainer.world_size))
88
 
89
  def training_step(self, batch, batch_idx):
90
  lr_dict = self.trainer.optimizers[0].next_batch()