Britania commited on
Commit
bcf45cf
·
verified ·
1 Parent(s): f039182

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +1 -1
handler.py CHANGED
@@ -17,7 +17,7 @@ class EndpointHandler():
17
  self.my_model_name,
18
  trust_remote_code=True,
19
  torch_dtype=torch.bfloat16,
20
- device_map="auto"
21
  )
22
  self.my_tokenizer = AutoTokenizer.from_pretrained(self.my_model_name, trust_remote_code=True)
23
 
 
17
  self.my_model_name,
18
  trust_remote_code=True,
19
  torch_dtype=torch.bfloat16,
20
+ device_map="cuda:0"
21
  )
22
  self.my_tokenizer = AutoTokenizer.from_pretrained(self.my_model_name, trust_remote_code=True)
23