Update koalpaca.py (#8)

- Update koalpaca.py (ee299e59aa54e845319082bc51ca2e04b70cc231)


Co-authored-by: yongsang yoo <4n3mone@users.noreply.huggingface.co>

koalpaca.py

@@ -8,6 +8,7 @@ class KoAlpaca(Model):
     def __init__(self):
         peft_model_id = "4n3mone/Komuchat-koalpaca-polyglot-12.8B"
         config = PeftConfig.from_pretrained(peft_model_id)
+        accelerator = Accelerator()
         self.bnb_config = BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_use_double_quant=True,
@@ -29,7 +30,7 @@ class KoAlpaca(Model):
                 inputs, 
                 return_tensors='pt', 
                 return_token_type_ids=False
-            ).to('cpu'), 
+            ).to(accelerator.device), 
             generation_config=self.gen_config
         )
         outputs = self.tokenizer.decode(output_ids[0]).split("### 답변: ")[-1]