added memory clearing

cisco-open · Apr 1, 2024 · 7769ebc · 7769ebc
1 parent 221d7f0
commit 7769ebc
Showing 1 changed file with 5 additions and 0 deletions.
diff --git 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}