Spaces:

xiang-wuu
/

yolov5

Runtime error

glenn-jocher commited on Jun 16, 2020

Commit

8db51c7

1 Parent(s): 4052603

tb_writer bug fix

Files changed (1) hide show

train.py CHANGED Viewed

@@ -195,8 +195,9 @@ def train(hyp):
     c = torch.tensor(labels[:, 0])  # classes
     # cf = torch.bincount(c.long(), minlength=nc) + 1.
     # model._initialize_biases(cf.to(device))
-    plot_labels(labels)
-    tb_writer.add_histogram('classes', c, 0)
     # Check anchors
     check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz)

     c = torch.tensor(labels[:, 0])  # classes
     # cf = torch.bincount(c.long(), minlength=nc) + 1.
     # model._initialize_biases(cf.to(device))
+    if tb_writer:
+        plot_labels(labels)
+        tb_writer.add_histogram('classes', c, 0)
     # Check anchors
     check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz)