Merge pull request #22 from kilianovski/main

Remove global 'torch.set_grad_enabled(False)' from FaceBoxes
choyingw · Aug 21, 2022 · d34e4af · d34e4af
2 parents b79d600 + 3d04863
commit d34e4af
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/FaceBoxes/FaceBoxes.py b/FaceBoxes/FaceBoxes.py
@@ -46,11 +46,13 @@ def viz_bbox(img, dets, wfp='out.jpg'):
 
 class FaceBoxes:
     def __init__(self, timer_flag=False):
-        torch.set_grad_enabled(False)
-
         net = FaceBoxesNet(phase='test', size=None, num_classes=2)  # initialize detector
         self.net = load_model(net, pretrained_path=pretrained_path, load_to_cpu=True)
         self.net.eval()
+
+        for p in self.net.parameters():
+            p.requires_grad_(False)
+
         # print('Finished loading model!')
 
         self.timer_flag = timer_flag