Merge pull request #51 from KPLauritzen/b-ensure-weights-dir-when-training

Ensures weights/ dir exists
5 years ago · 19e68e8a7b
parent 099e6f5ebd c3d4d321d3
commit 19e68e8a7b
1 changed files with 1 additions and 0 deletions
--- a/train.py
+++ b/train.py
@ -20,6 +20,7 @@ except:
    mixed_precision = False  # not installed

 wdir = 'weights' + os.sep  # weights dir
+os.makedirs(wdir, exist_ok=True)
 last = wdir + 'last.pt'
 best = wdir + 'best.pt'
 results_file = 'results.txt'