Update app.py

app.py

@@ -18,7 +18,6 @@ model_name = "sshleifer/distill-pegasus-cnn-16-4"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 tokenizer = PegasusTokenizer.from_pretrained(model_name, max_length=500)
 model = PegasusForConditionalGeneration.from_pretrained(model_name).to(device)
-# summarizer = pipeline(task="summarization", model=model_name, tokenizer=tokenizer, batch_size=1, device=-1)
 
 def summarize(slides):
     generated_slides = copy.deepcopy(slides)
@@ -26,8 +25,7 @@ def summarize(slides):
         for idx, (tag, content) in enumerate(contents):
             if tag.startswith('p'): 
                 try:
-                    # summary = summarizer(content, max_length=500)
-                    input = tokenizer(content, truncation=True, padding="longest", return_tensors="pt")
+                    input = tokenizer(content, truncation=True, padding="longest", return_tensors="pt").to(device)
                     tensor = model.generate(**input)
                     summary = tokenizer.batch_decode(tensor, skip_special_tokens=True)[0]
                     contents[idx] = (tag, summary)