Update app.py

app.py

@@ -24,10 +24,11 @@ class BertClass(torch.nn.Module):
     return output
 
 # Define models to be used
+device = "cuda" if torch.cuda.is_available() else "cpu"
 bert_path = "bert-base-uncased"
 bert_tokenizer = AutoTokenizer.from_pretrained(bert_path)
 bert_model = BertForSequenceClassification.from_pretrained(bert_path, num_labels=6)
-tuned_model = model = torch.load("pytorch_bert_toxic.bin", map_location=torch.device("cpu"))
+tuned_model = model = torch.load("pytorch_bert_toxic.bin", map_location=torch.device(device))
 
 # Read and format data.
 tweets_raw = pd.read_csv("test.csv", nrows=20)