Luis Chaves commited on
Commit
4c91d01
·
1 Parent(s): d86a1f5

use mps if available for mac

Browse files
Files changed (1) hide show
  1. src/everycure/extractor.py +6 -1
src/everycure/extractor.py CHANGED
@@ -33,7 +33,12 @@ MEDICAL_LABELS = [
33
  ]
34
 
35
  # Check for GPU availability
36
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
 
 
 
37
  logger.info(f"Using device: {device}")
38
 
39
  # Initialize model
 
33
  ]
34
 
35
  # Check for GPU availability
36
+ if torch.backends.mps.is_available():
37
+ device = torch.device("mps")
38
+ elif torch.cuda.is_available():
39
+ device = torch.device("cuda")
40
+ else:
41
+ device = torch.device("cpu")
42
  logger.info(f"Using device: {device}")
43
 
44
  # Initialize model