f15hb0wn commited on
Commit
0a79ce5
1 Parent(s): 37d0877

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -4
app.py CHANGED
@@ -5,14 +5,14 @@ import gradio as gr
5
  model_id = "witfoo/witq-1.0"
6
  dtype = torch.float16 # float16 for Tesla T4, V100, bfloat16 for Ampere+
7
  tokenizer = AutoTokenizer.from_pretrained(model_id)
 
 
8
  model = AutoModelForCausalLM.from_pretrained(
9
  model_id,
10
  torch_dtype=dtype,
11
- device_map="auto",
12
  )
13
- # Use GPU if available
14
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
15
- model.to(device)
16
 
17
  preamble = "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request."
18
 
 
5
  model_id = "witfoo/witq-1.0"
6
  dtype = torch.float16 # float16 for Tesla T4, V100, bfloat16 for Ampere+
7
  tokenizer = AutoTokenizer.from_pretrained(model_id)
8
+ # Use GPU if available
9
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
10
  model = AutoModelForCausalLM.from_pretrained(
11
  model_id,
12
  torch_dtype=dtype,
13
+ device_map=device,
14
  )
15
+
 
 
16
 
17
  preamble = "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request."
18