From 6e2b1bfe21d4c7444f5834fb9c2b50e7cee3efbc Mon Sep 17 00:00:00 2001
From: Johnson yue <Johnson-yue@users.noreply.github.com>
Date: Mon, 21 Dec 2020 18:04:07 +0800
Subject: [PATCH] add torch.no_grad() for inference network

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

diff --git a/main.py b/main.py
index e7aeb84..fd0b3a0 100644
--- a/main.py
+++ b/main.py
@@ -44,7 +44,8 @@ def detect(model, image, threshold=0.4, nms_iou=0.5):
     if HAS_CUDA:
         torch_image = torch_image.cuda()
 
-    hm, box, landmark = model(torch_image)
+    with torch.no_grad():
+        hm, box, landmark = model(torch_image)
     hm_pool = F.max_pool2d(hm, 3, 1, 1)
     scores, indices = ((hm == hm_pool).float() * hm).view(1, -1).cpu().topk(1000)
     hm_height, hm_width = hm.shape[2:]
@@ -132,4 +133,4 @@ def camera_demo():
     
 
 
-    
\ No newline at end of file
+