Bugfix : Prompt Inference 실행시, 모델의 device setting이 flag와 맞지 않는 버그 해결

[Se-Hun]

버그 발생 상황

GPU 환경에서 기존의 코드를 동작시킬 때 아래와 같은 에러가 발생합니다.

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument index in method wrapper__index_select)

해결 방법

위의 에러가 발생하는 원인은 KoGPTInference 클래스의 생성자가 호출될 때 입력 받은 device 변수를 통해 GPT 모델의 device가 설정되지 않아 발생하는 문제로 보입니다. 따라서, 아래와 같은 코드를 추가하여 문제를 해결하였습니다.

class KoGPTInference:
    def __init__(
            self,
            pretrained_model_name_or_path: Optional[Union[str, os.PathLike]],
            revision: str = 'KoGPT6B-ryan1.5b-float16',
            device: str = 'cuda',
            model_parallel: bool = False,
    ):

           ...

            self.model = GPTJForCausalLM.from_pretrained(
                  pretrained_model_name_or_path,  revision=revision,
                  pad_token_id=self.tokenizer.eos_token_id,
                  torch_dtype='auto', low_cpu_mem_usage=True
              ).to(device)

           ...