malvin noel commited on
Commit
339bd9a
·
1 Parent(s): 77f9296

modify cuda link

Browse files
Files changed (1) hide show
  1. scripts/generate_scripts.py +1 -1
scripts/generate_scripts.py CHANGED
@@ -15,7 +15,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
15
  model_id = "Qwen/Qwen2.5-0.5B"
16
 
17
  tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
18
- model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32, trust_remote_code=True)
19
 
20
 
21
  def generate_local(prompt: str, max_new_tokens: int = 350, temperature: float = 0.7) -> str:
 
15
  model_id = "Qwen/Qwen2.5-0.5B"
16
 
17
  tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
18
+ model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32, trust_remote_code=True).to(device)
19
 
20
 
21
  def generate_local(prompt: str, max_new_tokens: int = 350, temperature: float = 0.7) -> str: