보통 lr이 너무 커서 발생하는 문제일 수도 있고, 연산중에 분모에 0이 있거나 여러가지 이슈가 있는 것 같습니다. gpu만 사용했을때 발생하는거라 정확한 이유는 모르겠지만 아래의 세가지 글 한번 참고해서 파라미터 몇개 수정해보면 될 것 같습니다. 
- https://velog.io/@0hye/PyTorch-Nan-Loss-검출-방법
- https://data-newbie.tistory.com/281
- https://twinparadox.tistory.com/598

프레임워크를 tensorflow 사용했을 때 저랬던 경험 (train에는 정상적으로 뜨는데 valid는 nan이 뜨는 이슈)이 있는데요. 같은 코드인데 다른 tf버전에서 실행해보니 제대로 작동하던 경험이 있습니다.
버전도 한 번 점검해보시는 것을 추천드립니다.

loss에서 NaN이 나오는 경우는 크게 두가지입니다. 분모가 0이 되거나,  혹은 log안의 값이 0이 되면 NaN이 뜰 수 있습니다.