Update app.py

app.py

@@ -1,5 +1,6 @@
 import pickle
 import gradio as gr
+import torch
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from transformers import BertTokenizer, BertForSequenceClassification, pipeline, AutoTokenizer, AutoModelForSequenceClassification, TextClassificationPipeline, AutoModelForSeq2SeqLM, AutoModel, RobertaModel, RobertaTokenizer
 from sentence_transformers import SentenceTransformer
@@ -10,7 +11,7 @@ from fin_readability_sustainability import BERTClass, do_predict
 tokenizer_read = BertTokenizer.from_pretrained('ProsusAI/finbert')
 
 
-
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model_read = BERTClass(2, "readability")
 model_read.to(device)
 model_read.load_state_dict(torch.load('readability_model.bin', map_location=device)['model_state_dict'])