Update app.py

app.py

@@ -6,13 +6,14 @@ import torch
 
 
 class ModelProcessor:
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
     def __init__(self, repo_id="HuggingFaceTB/cosmo-1b"):
         # Initialize the tokenizer
         self.tokenizer = AutoTokenizer.from_pretrained(repo_id, use_fast=True)
 
         # Initialize and configure the model
         self.model = AutoModelForCausalLM.from_pretrained(
-            repo_id, torch_dtype=torch.float16, device_map="cpu", trust_remote_code=True
+            repo_id, torch_dtype=torch.float16, device_map={"": device}, trust_remote_code=True
         )
         self.model.eval()  # Set the model to evaluation mode