Fix a GPU memory leak in detect. No need to calculate gradients in inference. (#900)

Co-authored-by: JKO095 <juho-pekka.koponen@wartsila.com>
2022-10-08 00:50:04 +03:00 · 2022-10-08 00:50:04 +03:00 · 072f76c72c
parent 8b616af63a
commit 072f76c72c
1 changed files with 2 additions and 1 deletions
--- a/detect.py
+++ b/detect.py
@ -84,6 +84,7 @@ def detect(save_img=False):

        # Inference
        t1 = time_synchronized()
+        with torch.no_grad():   # Calculating gradients would cause a GPU memory leak
            pred = model(img, augment=opt.augment)[0]
        t2 = time_synchronized()