When I wrap my model in nn.DataParallel, it requires me to move the model to the GPU (via .cuda()). However, it also requires moving the inputs to the forward pass to CUDA (e.g., via to(torch.device('cuda'))). I saw this post 20 from 2017 mentioning DataParallel allows CPU inputs, but I’m running into issues when passing CPU tensors to the parallelized model (specifically, it’s saying it’s expecting cuda tensors but didn’t get cuda tensors). Maybe things changed since then.
I want to avoid a case where I put my input on one GPU, then DataParallel has to take it off that GPU and distribute it on the rest, making it really slow.
Are there any optimizations I can do to ensure that I’m not doing any unnecessary transfer between GPU and CPU? And is it correct that I have to pass cuda tensors to a parallelized module?
Thanks!