ElPlaguister commited on
Commit
46b206f
·
1 Parent(s): 80fb1ef

Fix KoAlpaca CUDA index

Browse files
Files changed (1) hide show
  1. koalpaca.py +1 -2
koalpaca.py CHANGED
@@ -16,8 +16,7 @@ class KoAlpaca(Model):
16
  )
17
  #self.model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, quantization_config=self.bnb_config, device_map={"":0})
18
  self.model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, quantization_config=self.bnb_config, device_map='auto')
19
- print(self.model.device)
20
- self.model = tp.tensor_parallel(self.model, [*next(self.model.parameters()).device])
21
  self.model = PeftModel.from_pretrained(self.model, peft_model_id)
22
  self.tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
23
  self.gen_config = GenerationConfig.from_pretrained('./models/koalpaca', 'gen_config.json')
 
16
  )
17
  #self.model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, quantization_config=self.bnb_config, device_map={"":0})
18
  self.model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, quantization_config=self.bnb_config, device_map='auto')
19
+ self.model = tp.tensor_parallel(self.model, ["cuda:0", "cuda:1"])
 
20
  self.model = PeftModel.from_pretrained(self.model, peft_model_id)
21
  self.tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
22
  self.gen_config = GenerationConfig.from_pretrained('./models/koalpaca', 'gen_config.json')