Update app.py

app.py

@@ -8,12 +8,14 @@ import spaces
 peft_model_id = "rootxhacker/CodeAstra-7B"
 config = PeftConfig.from_pretrained(peft_model_id)
 model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, return_dict=True, load_in_4bit=True, device_map='auto')
+model.to('cuda')
 tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
 
 # Load the Lora model
 model = PeftModel.from_pretrained(model, peft_model_id)
 
 
+
 @spaces.GPU(duration=200)
 def get_completion(query, model, tokenizer):
     inputs = tokenizer(query, return_tensors="pt")