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

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -0
app.py CHANGED
@@ -10,6 +10,9 @@ model = AutoModelForCausalLM.from_pretrained(
10
  torch_dtype=dtype,
11
  device_map="auto",
12
  )
 
 
 
13
 
14
  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."
15
 
 
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