YifengDing commited on
Commit
f5b7d69
1 Parent(s): 7447bf0
Files changed (1) hide show
  1. app.py +2 -3
app.py CHANGED
@@ -8,11 +8,10 @@ import gradio as gr
8
 
9
 
10
  def main(
11
- base_model="ise-uiuc/Magicoder-S-DS-6.7B",
12
- device="cuda" if torch.cuda.is_available() else "cpu",
13
  ):
14
  pipeline = transformers.pipeline(
15
- "text-generation", model=base_model, torch_dtype=torch.bfloat16, device=device
16
  )
17
 
18
  def evaluate_magicoder(
 
8
 
9
 
10
  def main(
11
+ base_model="ise-uiuc/Magicoder-S-DS-6.7B"
 
12
  ):
13
  pipeline = transformers.pipeline(
14
+ "text-generation", model=base_model, torch_dtype=torch.bfloat16, device_map="auto"
15
  )
16
 
17
  def evaluate_magicoder(