From 072f76c72c641c7a1ee482e39f604f6f8ef7ee92 Mon Sep 17 00:00:00 2001 From: jpkoponen <60470320+jpkoponen@users.noreply.github.com> Date: Sat, 8 Oct 2022 00:50:04 +0300 Subject: [PATCH] Fix a GPU memory leak in detect. No need to calculate gradients in inference. (#900) Co-authored-by: JKO095 --- detect.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/detect.py b/detect.py index 5648a91b32..5e0c4416a4 100644 --- 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