tb_writer bug fix

5 years ago · 8db51c7002
parent 4052603e1f
commit 8db51c7002
1 changed files with 3 additions and 2 deletions
--- a/train.py
+++ b/train.py
@ -195,6 +195,7 @@ def train(hyp):
    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)