updates

test.py

@@ -44,8 +44,6 @@ def test(
     outputs, mAPs, mR, mP, TP, confidence, pred_class, target_class = [], [], [], [], [], [], [], []
     AP_accum, AP_accum_count = np.zeros(nC), np.zeros(nC)
     for batch_i, (imgs, targets) in enumerate(dataloader):
-
-        with torch.no_grad():
         output = model(imgs.to(device))
         output = non_max_suppression(output, conf_thres=conf_thres, nms_thres=nms_thres)
 
@@ -134,6 +132,7 @@ if __name__ == '__main__':
     opt = parser.parse_args()
     print(opt, end='\n\n')
 
+    with torch.no_grad():
         mAP = test(
             opt.cfg,
             opt.data_cfg,

train.py

@@ -195,6 +195,7 @@ def train(
             os.system('cp ' + latest + ' ' + os.path.join(weights, 'backup{}.pt'.format(epoch)))
 
         # Calculate mAP
+        with torch.no_grad():
             mAP, R, P = test.test(cfg, data_cfg, weights=latest, batch_size=batch_size, img_size=img_size)
 
         # Write epoch results