Update app.py

app.py

@@ -7,6 +7,12 @@ from peft import AutoPeftModelForCausalLM
 from transformers import AutoTokenizer
 import torch
 
+import torch
+print(f"Is CUDA available: {torch.cuda.is_available()}")
+# True
+print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
+# Tesla T4
+
 @spaces.GPU
 def main():
 
@@ -25,6 +31,8 @@ def main():
         messages = [{"role": "user", "content": prompts}]
         
         input_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")
+        input_ids = input_ids.to("cuda")
+        
         prompt_padded_len = len(input_ids[0])
         
         # Generate the translation
@@ -63,6 +71,7 @@ def main():
 
     
     model = AutoPeftModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, use_auth_token=hf_token)
+    model = model.to("cuda")
     tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=hf_token)
     
     tokenizer_with_prefix_space = AutoTokenizer.from_pretrained(model_id, add_prefix_space=True, use_auth_token=hf_token)