add torch.no_grad()

clovaai · Dec 4, 2019 · e332dd8 · e332dd8
1 parent 3cd65f5
commit e332dd8
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/test.py b/test.py
@@ -81,15 +81,17 @@ def test_net(net, image, text_threshold, link_threshold, low_text, cuda, poly, r
         x = x.cuda()
 
     # forward pass
-    y, feature = net(x)
+    with torch.no_grad():
+        y, feature = net(x)
 
     # make score and link map
     score_text = y[0,:,:,0].cpu().data.numpy()
     score_link = y[0,:,:,1].cpu().data.numpy()
 
     # refine link
     if refine_net is not None:
-        y_refiner = refine_net(y, feature)
+        with torch.no_grad():
+            y_refiner = refine_net(y, feature)
         score_link = y_refiner[0,:,:,0].cpu().data.numpy()
 
     t0 = time.time() - t0