elghafiani commited on
Commit
726ccef
·
1 Parent(s): 8ba49f1

change device_map="auto"

Browse files
Files changed (1) hide show
  1. app.py +1 -2
app.py CHANGED
@@ -9,7 +9,6 @@ import gradio as gr
9
 
10
  def main(
11
  base_model="ise-uiuc/Magicoder-S-DS-6.7B",
12
- device="cuda:0",
13
  port=8080,
14
  ):
15
  tokenizer = AutoTokenizer.from_pretrained(base_model)
@@ -17,7 +16,7 @@ def main(
17
  "text-generation",
18
  model=base_model,
19
  torch_dtype=torch.float16,
20
- device=device
21
  )
22
  def evaluate_magicoder(
23
  instruction,
 
9
 
10
  def main(
11
  base_model="ise-uiuc/Magicoder-S-DS-6.7B",
 
12
  port=8080,
13
  ):
14
  tokenizer = AutoTokenizer.from_pretrained(base_model)
 
16
  "text-generation",
17
  model=base_model,
18
  torch_dtype=torch.float16,
19
+ device_map="auto"
20
  )
21
  def evaluate_magicoder(
22
  instruction,