Update app.py

app.py

@@ -4,8 +4,6 @@ import os
 import torch
 import spaces
 
-device = "cuda:0" if torch.cuda.is_available() else "cpu"
-print(f'[INFO] Using device: {device}')
 
 # Define model paths
 MODEL_PATHS = {
@@ -20,7 +18,9 @@ TOKEN = os.environ['TOKEN']
 
 # Translation function for Nano and Large models
 @spaces.GPU
-def translate_nano_large(text, model_path, device='cuda:0'):
+def translate_nano_large(text, model_path):
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    print(f'[INFO] Using device: {device}')
     translator = pipeline("translation", model=model_path, token=TOKEN)
     translated = translator(
         text, 
@@ -38,7 +38,9 @@ def translate_nano_large(text, model_path, device='cuda:0'):
 
 # Translation function for Ultra and Supreme models
 @spaces.GPU
-def translate_ultra_supreme(text, model_path, device='cuda:0'):
+def translate_ultra_supreme(text, model_path):
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    print(f'[INFO] Using device: {device}')
     model = AutoModelForSeq2SeqLM.from_pretrained(model_path, token=TOKEN)
     tokenizer = AutoTokenizer.from_pretrained(model_path, src_lang="eng_Latn", tgt_lang="ary_Arab", token=TOKEN)
     translator = pipeline(