alonsosilva commited on
Commit
e287122
·
unverified ·
1 Parent(s): ca2a338

Add autoselect device for mps

Browse files
Files changed (1) hide show
  1. app.py +7 -1
app.py CHANGED
@@ -8,7 +8,13 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStream
8
  from transformers.generation import LogitsProcessor
9
  from typing_extensions import TypedDict
10
 
11
- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 
 
 
 
 
12
 
13
  model_id = "Qwen/Qwen3-0.6B"
14
  model = AutoModelForCausalLM.from_pretrained(model_id).to(device)
 
8
  from transformers.generation import LogitsProcessor
9
  from typing_extensions import TypedDict
10
 
11
+ # Auto select device (CUDA > MPS > CPU)
12
+ if torch.cuda.is_available():
13
+ device = torch.device("cuda")
14
+ elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
15
+ device = torch.device("mps")
16
+ else:
17
+ device = torch.device("cpu")
18
 
19
  model_id = "Qwen/Qwen3-0.6B"
20
  model = AutoModelForCausalLM.from_pretrained(model_id).to(device)