Update koalpaca.py

koalpaca.py

@@ -7,6 +7,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,
@@ -28,7 +29,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]