Update app.py

app.py

@@ -5,16 +5,15 @@ import torch
 
 os.environ["TRANSFORMERS_CACHE"] = "/tmp"
 
-# Suppress NVML warning
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"
-torch.cuda.init()
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+if torch.cuda.is_available():
+    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+    torch.cuda.init()
 
 MODEL_NAME = "s-nlp/roberta-base-formality-ranker"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
 
-# Move model to GPU if available
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = model.to(device)
 
 def calculate_formality_percentages(score):
@@ -68,5 +67,5 @@ if __name__ == "__main__":
         server_name="0.0.0.0",
         server_port=7860,
         share=True, 
-        enable_queue=True 
+        enable_queue=True  
     )