mirror of
https://github.com/WongKinYiu/yolov7.git
synced 2024-11-21 01:11:37 +08:00
Fix a GPU memory leak in detect. No need to calculate gradients in inference. (#900)
Co-authored-by: JKO095 <juho-pekka.koponen@wartsila.com>
This commit is contained in:
parent
8b616af63a
commit
072f76c72c
@ -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
|
||||
|
Loading…
Reference in New Issue
Block a user