Sounds like this error is expected here.
The input encoder you posted earlier will always run on CPU, since you don’t pass a device kwarg to torch.zeroes. I’m assuming you’re calling the encoder from within the forward function. If you want inputs to be distributed to all GPUs, you need to call the wrapped module (the resulting model after wrapping it with nn.DataParallel) with the CPU side inputs, and nn.DataParallel will make sure the inputs are distributed accordingly. If you generate the encoded input from within the forward function, there is no place where nn.DataParallel could hook into and move them around.