token

app.py

@@ -22,7 +22,7 @@ base_model = AutoModelForCausalLM.from_pretrained(
     trust_remote_code=True,
     token=huggingface_token,  # Use the token parameter
     cache_dir=cache_dir  # Specify cache directory
-).to("cuda")  # Move model to CUDA
+)
 
 # Load the tokenizer
 tokenizer = AutoTokenizer.from_pretrained(
@@ -34,7 +34,7 @@ tokenizer = AutoTokenizer.from_pretrained(
 )
 
 # Load the fine-tuned model
-ft_model = PeftModel.from_pretrained(base_model, "checkpoint-2800", cache_dir=cache_dir).to("cuda")  # Move model to CUDA
+ft_model = PeftModel.from_pretrained(base_model, "checkpoint-2800", cache_dir=cache_dir)
 
 def formatting_func(job_description):
     text = f"### The job description: {job_description}\n ### The skills: "
@@ -47,6 +47,8 @@ def generate_skills(job_description):
 
     ft_model.eval()
     with torch.no_grad():
+        # Move model to GPU if not already there
+        ft_model.to("cuda")
         output_tokens = ft_model.generate(**model_input, max_new_tokens=200)[0]
         generated_text = tokenizer.decode(output_tokens, skip_special_tokens=True)