[PyTorch] DataLoader shuffle 기능 사용 시, RuntimeError: Expected a 'cuda' device type for generator but found 'cpu'

PyTorch에서 dataloader 기능 중에 shuffle을 사용하고자했는데 다음과 같은 에러를 봤다.

 

RuntimeError: Expected a 'cuda' device type for generator but found 'cpu'

 

참고로 나는 내가 만든 custom dataset을 class로 만들고, 그 클래스에 대하여 DataLoader를 만들었다.

 

인터넷을 검색하여 찾아보니 이 shuffle 기능을 위해 사용하는 generator가 cpu로 설정되어 있으면 그 device에서 generator를 생성하는 것으로 보였다.

나는 custom dataset에서 모두 데이터의 device를 cuda로 설정하는 작업을 거치기 때문에 오류가 생긴 것이다.

 

 

이 글에서 보면

generator를 설정할 수 있다.

 

따라서 

device = "cuda"
test_ds = CustomDataset_test()
test_dl = DataLoader(test_ds, batch_size=16, shuffle=True, generator=torch.Generator(device=device) )

위의 코드처럼 generator의 device를 cuda로 설정하면 문제가 발생하지 않는다.

 

문제 해결!

 

---

 

참고자료

 

https://pytorch.org/docs/stable/generated/torch.Generator.html#torch.Generator

Generator — PyTorch 1.12 documentation