From 8db51c70025fd28ee7bd7ca60b41f183bab60c0e Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Tue, 16 Jun 2020 16:05:28 -0700
Subject: [PATCH] tb_writer bug fix

---
 train.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/train.py b/train.py
index e9be7e4..7df99c0 100644
--- a/train.py
+++ b/train.py
@@ -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)
+    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)