Kvikontent commited on
Commit
acf9e66
1 Parent(s): e779eca

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -2,7 +2,7 @@ import streamlit as st
2
  from transformers import AutoTokenizer, FalconModel
3
  import torch
4
 
5
- device = "cpu"
6
  tokenizer = AutoTokenizer.from_pretrained("Rocketknight1/falcon-rw-1b")
7
  model = FalconModel.from_pretrained("Rocketknight1/falcon-rw-1b")
8
 
 
2
  from transformers import AutoTokenizer, FalconModel
3
  import torch
4
 
5
+ device = "cuda" if torch.cuda.is_available() else "cpu"
6
  tokenizer = AutoTokenizer.from_pretrained("Rocketknight1/falcon-rw-1b")
7
  model = FalconModel.from_pretrained("Rocketknight1/falcon-rw-1b")
8