Update app.py

app.py

@@ -17,15 +17,13 @@ ocr_model_name = "PleIAs/OCRonos-Vintage"
 import torch
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
 
-device = "cuda"
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Load pre-trained model and tokenizer
 model_name = "PleIAs/OCRonos-Vintage"
 model = GPT2LMHeadModel.from_pretrained(model_name)
 tokenizer = GPT2Tokenizer.from_pretrained(model_name)
 
-model.to(device)
-
 # CSS for formatting
 css = """
 <style>
@@ -166,6 +164,7 @@ def split_text(text, max_tokens=500):
 # Function to generate text
 @spaces.GPU
 def ocr_correction(prompt, max_new_tokens=500):
+    model.to(device)
     
     prompt = f"""### Text ###\n{prompt}\n\n\n### Correction ###\n"""
     input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)