Implement gc.collect() to free memory (#10129)

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
This commit is contained in:
Glenn Jocher 2024-04-17 19:37:40 -07:00 committed by GitHub
parent 37ffebd690
commit 1f4bed233a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 7 additions and 0 deletions

View file

@ -1,5 +1,6 @@
# Ultralytics YOLO 🚀, AGPL-3.0 license
import gc
import math
import os
import random
@ -581,6 +582,7 @@ def profile(input, ops, n=10, device=None):
except Exception as e:
LOGGER.info(e)
results.append(None)
gc.collect() # attempt to free unused memory
torch.cuda.empty_cache()
return results