From 7769ebc387275afb134073e86d235b08157916e9 Mon Sep 17 00:00:00 2001 From: unknown Date: Mon, 1 Apr 2024 17:18:33 +0400 Subject: [PATCH] added memory clearing --- Ethosight/Ethosight/EthosightCore.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/Ethosight/Ethosight/EthosightCore.py b/Ethosight/Ethosight/EthosightCore.py index cd80abf..e61e546 100644 --- a/Ethosight/Ethosight/EthosightCore.py +++ b/Ethosight/Ethosight/EthosightCore.py @@ -275,6 +275,11 @@ def compute_affinity_scores(self, label_to_embeddings, image_path, normalize_fn= top_labels = sorted_labels[:100] top_scores = sorted_scores[:100] + # Clear memory + del vision_embeddings, raw_scores, unique_label_embeddings, inputs + if torch.cuda.is_available(): + torch.cuda.empty_cache() # Clear unused memory + # Return the labels and scores as a dictionary return {'labels': top_labels, 'scores': top_scores}