W&B ID reset on training completion (#1852)

2021-01-06 16:39:03 -08:00
parent 84ad6080ae
commit 4f2341c0ad
1 changed files with 2 additions and 2 deletions
@@ -350,8 +350,8 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non
 def strip_optimizer(f='weights/best.pt', s=''):  # from utils.general import *; strip_optimizer()
    # Strip optimizer from 'f' to finalize training, optionally save as 's'
    x = torch.load(f, map_location=torch.device('cpu'))
-    x['optimizer'] = None
-    x['training_results'] = None
+    for key in 'optimizer', 'training_results', 'wandb_id':
+        x[key] = None
    x['epoch'] = -1
    x['model'].half()  # to FP16
    for p in x['model'].parameters():