Kunal Pai commited on
Commit
ab59793
·
1 Parent(s): 941e057

MPS support for memories for Mac users

Browse files
Files changed (1) hide show
  1. src/manager/manager.py +6 -1
src/manager/manager.py CHANGED
@@ -164,7 +164,12 @@ class GeminiManager:
164
  return []
165
  top_k = min(k, len(memories))
166
  # Semantic Retrieval with GPU
167
- device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 
 
 
 
168
  print(f"Using device: {device}")
169
  model = SentenceTransformer('all-MiniLM-L6-v2', device=device)
170
  doc_embeddings = model.encode(memories, convert_to_tensor=True, device=device)
 
164
  return []
165
  top_k = min(k, len(memories))
166
  # Semantic Retrieval with GPU
167
+ if torch.cuda.is_available():
168
+ device = 'cuda'
169
+ elif torch.backends.mps.is_available() and torch.backends.mps.is_built():
170
+ device = 'mps'
171
+ else:
172
+ device = 'cpu'
173
  print(f"Using device: {device}")
174
  model = SentenceTransformer('all-MiniLM-L6-v2', device=device)
175
  doc_embeddings = model.encode(memories, convert_to_tensor=True, device=device)