I met an interesting (and breaking) behaviour with DataParallel when I override forward() by assignment. Something like:
class MyEmbedding(nn.Embedding):
    def __init__(self, swap_fwd=True, *args, **kwargs):
        super().__init__(*args, **kwargs)
        if swap_fwd:
            self.FC = nn.Linear(self.embedding_dim, self.embedding_dim)
            self.forward = self.forward_
        else:
            self.FC = None

    def forward_(self, inputs):
        embeds = super().forward(inputs)
        return self.FC(embeds)

This works fine on a single gpu, but with multiple, I get “arguments are located on different GPUs”. It seems that when _forward() is executing, self is not the appropriate replica (on the correct gpu), but is the original object (typically on gpu 0).
If instead I override forward() in the normal way, everything works fine. Something like:
...
    def forward(self, inputs):
        embeds = super().forward(inputs)
        if self.FC is not None:
            return self.FC(embeds)
        return embeds

I’d be interested to know what’s going on?