cpu mode

models/distilbart_cnn_12_6.py

@@ -4,7 +4,7 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 # loading the model outside of the function makes it faster
 SUMMARIZATION_MODEL = "sshleifer/distilbart-cnn-12-6"
 tokenizer   = AutoTokenizer.from_pretrained(SUMMARIZATION_MODEL)
-model       = AutoModelForSeq2SeqLM.from_pretrained(SUMMARIZATION_MODEL, device_map="cuda:0")
+model       = AutoModelForSeq2SeqLM.from_pretrained(SUMMARIZATION_MODEL)
 
 def summarize(text, max_len=20):
     """

models/t5_small_medium_title_generation.py

@@ -4,7 +4,7 @@ import torch
 
 def t5model(prompt: str) -> str:
     tokenizer = AutoTokenizer.from_pretrained("fabiochiu/t5-small-medium-title-generation")
-    model = AutoModelForSeq2SeqLM.from_pretrained("fabiochiu/t5-small-medium-title-generation", device_map="cuda:0", torch_dtype=torch.float16)
+    model = AutoModelForSeq2SeqLM.from_pretrained("fabiochiu/t5-small-medium-title-generation", torch_dtype=torch.float16)
     inputs = tokenizer(
         ["summarize:" + prompt],
         return_tensors="pt",