Issue while running the model over multiple GPUs using nn.DataParallel

[ShoRit]

I used the nn.DataParallel() method in an attempt to do multi-processing. However, I am always met with this error.

RuntimeError: Expected tensor for argument #1 'input' to have the same device as tensor for argument #2 'weight'; but device 1 does not equal 0 (while checking arguments for cudnn_batch_norm) in the line.

head = self.bn0(head)

Is there a way to remedy this? How did you carry out Multi-processing?

[ShoRit]

One thing I discovered is that the tensors in self.embedding.weight are always on the device 0, as opposed to the data which have been split. This raises the question whether the model is correctly being copied onto the 4 devices.

head cuda:3, embedding_weight cuda:0
head cuda:2, embedding_weight cuda:0
head cuda:1, embedding_weight cuda:0
head cuda:0, embedding_weight cuda:0

[apoorvumang]

Unfortunately multi-GPU is not yet supported in this code by default, you will have to make modifications. I would recommend using huggingface accelerate.