[huggingface] not enough values to unpack (expected 2, got 1)

황준하·2024년 7월 6일

허깅페이스 모델(Encoder-only LM)을 이용해 inference를 했을 때 만난 오류

테스트용으로 output의 형태를 보기 위해 시도한 것이었고 데이터(텍스트)를 1개만 사용했었다.
하지만, huggingface model의 input으로는 2D tensor가 들어올 것을 가정하기 때문에 발생한 오류였다.

따라서, 아래와 같이 .unsqueeze(0) 을 붙여 2D로 만들어주니 정상적으로 작동했다.

koelectra.to('cuda:0')

out = koelectra(input_ids = train_dataset.__getitem__(970)['input_ids'].unsqueeze(0).to('cuda:0'),
          attention_mask = train_dataset.__getitem__(970)['attention_mask'].unsqueeze(0).to('cuda:0'),
          token_type_ids = train_dataset.__getitem__(970)['token_type_ids'].unsqueeze(0).to('cuda:0'),
         )

out

0개의 댓글