huggingface trainer api를 colab gpu와 함께 사용하는 방법

평화롭게 huggingface 모듈로 bert 모델을 fine tuning하기 위해 학습시키던 도중

훈련도 돠 됬겠다 기쁜 마음으로 pipeline에 적용하려고 했으나, 요런 에러가 등장했다. 

 

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

 

 

처음에는 아 model.to(device)를 깜빡했네! 라고 생각해서 바꿔줬는데

아무리 trainer 내부의 모델을 to(device) 처리해줘도 아무런 변화가없었다...

 

결론부터 말하자면, 원인을 해결하는 문제는 아직 찾지 못했고, 대안만 찾았다.

 

구글링으로는 별다른 대안을 찾지 못했고, 

input data 역시 to(device) 처리를 해줘야 하는건가? 라는 생각에 dataset에 적용했으나, dataset에는 to(device)를 적용할 수 없는듯했다.

 

그 외에도 여러 방법을 시도했으나 다 실패해서, 그냥 모델을 저장하고 다시 불러오기로 타협을 했더니, 파이프라인이 잘 작동했다.

trainer.save_model("./bert")	#모델 저장.
model = AutoModelForMaskedLM.from_pretrained("./bert")	#저장한모델 불러오기 

from transformers import pipeline
mask_filler = pipeline(
    "fill-mask", model=model, 
    tokenizer=tokenizer
)

 

 

혹시 에러 원인을 아시는 분은 댓글 남겨주시면 감사합니다!