to device

app.py

@@ -4,7 +4,7 @@ import re
 
 model = None
 tokenizer = None
-device = "cpu"
+# device = "cpu"
 
 def init():
     from transformers import MT5ForConditionalGeneration, T5TokenizerFast
@@ -13,8 +13,8 @@ def init():
     global model, tokenizer
 
     hf_token = os.environ.get("HF_TOKEN")
-    model_path = "lm-spell/mt5-base-ft-ssc"
-    model = MT5ForConditionalGeneration.from_pretrained(model_path, token=hf_token).to(device)
+
+    model = MT5ForConditionalGeneration.from_pretrained("lm-spell/mt5-base-ft-ssc", token=hf_token)
     tokenizer = T5TokenizerFast.from_pretrained("google/mt5-base")
     tokenizer.add_special_tokens({'additional_special_tokens': ['<ZWJ>']})
         
@@ -29,7 +29,7 @@ def correct(text):
         padding='do_not_pad',
         max_length=1024
     )
-    inputs = {k: v.to(device) for k, v in inputs.items()}
+    # inputs = {k: v.to(device) for k, v in inputs.items()}
 
     with torch.inference_mode():
         outputs = model.generate(