@ggaemo your example looks simple enough, don’t see any immediate issues.
Something is really wrong though based on the fact the main process is pinning all of your cores in the kernel (red). Something is not happy at a driver/IRQ/IO level. Check to make sure CUDA / GPU drivers match and are compat with your PyTorch, maybe update all of them to latest to be sure. Make sure the drives where your dataset is located are working properly (try copying your dataset to another location and make sure that happens quickly with no issues). Try compiling/running some CUDA sample benchmark apps (the bandwidth tests in utilities perhaps).