Update app.py

app.py

@@ -7,9 +7,14 @@ from transformers import pipeline, T5Tokenizer, T5ForConditionalGeneration
 # https://huggingface.co/docs/hub/spaces-sdks-gradio
 # model = GPT2LMHeadModel.from_pretrained("binxu/Ziyue-GPT2-deep")
 # generator = pipeline('text-generation', model=model, tokenizer='bert-base-chinese')
+if torch.cuda.is_available():
+    device = 0
+else: 
+    device = -1
+
 tokenizer = T5Tokenizer.from_pretrained("Langboat/mengzi-t5-base")
 model = T5ForConditionalGeneration.from_pretrained("binxu/mengzi-t5-base-finetuned-punctuation")
-text2text_generator = pipeline("text2text-generation", model=model, tokenizer=tokenizer, device=0)
+text2text_generator = pipeline("text2text-generation", model=model, tokenizer=tokenizer, device=device)
 
 def generate(prompt, ):
     torch.manual_seed(42)