Update app.py

app.py

@@ -1,11 +1,16 @@
 import streamlit as st
 from transformers import pipeline
 
+import torch
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 st.header("Sentiment-demo-app")
 st.subheader("Please be patient and wait up to a minute until the demo app is loaded.")
 st.caption("This is a very simple demo application for a zero-shot classification pipeline to classify positive, neutral, or negative sentiment for a short text. Enter your text in the box below and press CTRl+ENTER to run the model.")
 
 classifier = pipeline("zero-shot-classification", model='facebook/bart-large-mnli')
+model.to(device)
 
 text = st.text_area('Enter text here!')
 candidate_labels = ['Positive', 'Neutral', 'Negative']