updates

train.py

@@ -321,18 +321,16 @@ def train():
         final_epoch = epoch + 1 == epochs
         if opt.prebias:
             print_model_biases(model)
-        else:
-            # Calculate mAP
-            if not opt.notest or final_epoch:
-                with torch.no_grad():
-                    results, maps = test.test(cfg,
-                                              data,
-                                              batch_size=batch_size,
-                                              img_size=opt.img_size,
-                                              model=model,
-                                              conf_thres=0.001 if final_epoch and epoch > 0 else 0.1,  # 0.1 for speed
-                                              save_json=final_epoch and epoch > 0 and 'coco.data' in data,
-                                              dataloader=testloader)
+        elif not opt.notest or final_epoch:  # Calculate mAP
+            with torch.no_grad():
+                results, maps = test.test(cfg,
+                                          data,
+                                          batch_size=batch_size,
+                                          img_size=opt.img_size,
+                                          model=model,
+                                          conf_thres=0.001 if final_epoch else 0.1,  # 0.1 for speed
+                                          save_json=final_epoch and 'coco.data' in data and model.nc == 80,
+                                          dataloader=testloader)
 
         # Write epoch results
         with open(results_file, 'a') as f: