Fix number of epochs at training

lianggong24 · Mar 14, 2022 · 312d24f · 312d24f
1 parent 2bbea4c
commit 312d24f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/bonito/training.py b/bonito/training.py
@@ -204,7 +204,7 @@ def fit(self, workdir, epochs=1, lr=2e-3, **optim_kwargs):
 
         lr_scheduler = self.get_lr_scheduler(epochs, last_epoch=last_epoch)
 
-        for epoch in range(1 + last_epoch, epochs + 1 + last_epoch):
+        for epoch in range(1 + last_epoch, epochs + 1):
             try:
                 with bonito.io.CSVLogger(os.path.join(workdir, 'losses_{}.csv'.format(epoch))) as loss_log:
                     train_loss, duration = self.train_one_epoch(loss_log, lr_scheduler)