테스트용으로 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