Did you use num_workers > 0 and did you use your model on the GPU?
If so, it might be a bit problematic, since the DataLoader wants to use multi-processing, which apparently clashes with the CUDA calls.
With num_workers=0 it’ll work, but I’m not sure, if it will be the bottleneck in your application, since the main thread will be used to load the data.
Another approach would be to store all processed features as a Tensor beforehand and then just load it in a Dataset.