From c2809622c1143eb9a67d97aba0b3214b20816c64 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Wed, 17 Apr 2019 17:35:00 +0200
Subject: [PATCH] updates

---
 train.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/train.py b/train.py
index df26a531..b5410b46 100644
--- a/train.py
+++ b/train.py
@@ -201,7 +201,9 @@ def train(
         # Calculate mAP
         opt.notest = opt.notest or (opt.nosave and epoch < 10)  # skip testing first 10 epochs if opt.nosave
         if not opt.notest or epoch == epochs - 1:  # always test final epoch
-            results = test.test(cfg, data_cfg, batch_size=batch_size, img_size=img_size, model=model, conf_thres=0.1)
+            with torch.no_grad():
+                results = test.test(cfg, data_cfg, batch_size=batch_size, img_size=img_size, model=model,
+                                    conf_thres=0.1)
 
         # Write epoch results
         with open('results.txt', 'a') as file: