Update functions.py

functions.py

@@ -46,7 +46,7 @@ def load_model(peft_model_id):
     config = PeftConfig.from_pretrained(peft_model_id)
     print(['load_model', 'loading model'])
     model = AutoModelForCausalLM.from_pretrained(
-        config.base_model_name_or_path, return_dict=True, load_in_8bit=True, device_map='cuda')
+        config.base_model_name_or_path, return_dict=True, load_in_8bit=True, device_map='balanced')
     print(['load_model', 'loading tokenizer'])
     tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
     model = PeftModel.from_pretrained(model, peft_model_id)