Fix a GPU memory leak in detect. No need to calculate gradients in in…

…ference. (#900) Co-authored-by: JKO095 <[email protected]>
WongKinYiu · Oct 7, 2022 · 072f76c · 072f76c
1 parent 8b616af
commit 072f76c
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/detect.py b/detect.py
@@ -84,7 +84,8 @@ def detect(save_img=False):
 
         # Inference
         t1 = time_synchronized()
-        pred = model(img, augment=opt.augment)[0]
+        with torch.no_grad():   # Calculating gradients would cause a GPU memory leak
+            pred = model(img, augment=opt.augment)[0]
         t2 = time_synchronized()
 
         # Apply NMS