Spaces:
Running
on
Zero
Running
on
Zero
gpu optims
Browse files
modular_graph_and_candidates.py
CHANGED
@@ -130,7 +130,7 @@ def embedding_similarity_clusters(models_root: Path, missing: List[str], thr: fl
|
|
130 |
all_embeddings = []
|
131 |
|
132 |
print(f"Encoding embeddings for {len(names)} models...")
|
133 |
-
batch_size =
|
134 |
|
135 |
for i in tqdm(range(0, len(names), batch_size), desc="Batches", leave=False):
|
136 |
batch_names = names[i:i+batch_size]
|
@@ -142,9 +142,10 @@ def embedding_similarity_clusters(models_root: Path, missing: List[str], thr: fl
|
|
142 |
all_embeddings.append(emb)
|
143 |
print(f"✓ Completed batch of {len(batch_names)} models")
|
144 |
|
145 |
-
# Clear GPU cache every
|
146 |
-
if i % (
|
147 |
torch.cuda.empty_cache()
|
|
|
148 |
print(f"🧹 Cleared GPU cache after batch {i//batch_size + 1}")
|
149 |
|
150 |
except Exception as e:
|
|
|
130 |
all_embeddings = []
|
131 |
|
132 |
print(f"Encoding embeddings for {len(names)} models...")
|
133 |
+
batch_size = 4 # Reduced to be more conservative
|
134 |
|
135 |
for i in tqdm(range(0, len(names), batch_size), desc="Batches", leave=False):
|
136 |
batch_names = names[i:i+batch_size]
|
|
|
142 |
all_embeddings.append(emb)
|
143 |
print(f"✓ Completed batch of {len(batch_names)} models")
|
144 |
|
145 |
+
# Clear GPU cache every 3 batches to prevent memory accumulation
|
146 |
+
if i % (3 * batch_size) == 0 and torch.cuda.is_available():
|
147 |
torch.cuda.empty_cache()
|
148 |
+
torch.cuda.synchronize() # Force GPU sync
|
149 |
print(f"🧹 Cleared GPU cache after batch {i//batch_size + 1}")
|
150 |
|
151 |
except Exception as e:
|