Fix img_size naming in hyperparameters logging

5 years ago · 38acc5f3c5
parent e558963244
commit 38acc5f3c5
1 changed files with 1 additions and 1 deletions
--- a/train.py
+++ b/train.py
@ -62,7 +62,7 @@ def train(hyp):
    if tb_writer:
        tb_hparams_dict = hyp
        tb_hparams_dict.update(vars(opt))
-        tb_hparams_dict['img_size_w'], tb_hparams_dict['img_size_h'] = tb_hparams_dict['img_size']
+        tb_hparams_dict['img_size_train'], tb_hparams_dict['img_size_test'] = tb_hparams_dict['img_size']
        del tb_hparams_dict['img_size']
        tb_writer.add_hparams(tb_hparams_dict, {})