there are model wrapped by nn.DataParallel
self.model = Bert(6, 12, 513, 384*4, 64, 64, 2, 384, self.base_task.max_vocab_indexes['input_ids'])
self.model = nn.DataParallel(self.model).cuda()

and in model there are constant tensor, which name is pos.
class Embedding(nn.Module):
    def __init__(self, maxlen, d_model, n_segments, vocab_size, device='cuda'):
        super(Embedding, self).__init__()
        self.device = device
        self.tok_embed = nn.Embedding(vocab_size, d_model)  # token embedding
        self.pos_embed = nn.Embedding(maxlen, d_model)  # position embedding
        self.seg_embed = nn.Embedding(n_segments, d_model)  # segment(token type) embedding
        self.norm = nn.LayerNorm(d_model)

    def forward(self, x, seg):
        seq_len = x.size(1)
        pos = torch.arange(seq_len, dtype=torch.long, device=self.device)
        pos = pos.unsqueeze(0).expand_as(x)  # (seq_len,) -> (batch_size, seq_len)
        embedding = self.tok_embed(x) + self.pos_embed(pos) + self.seg_embed(seg)
        return self.norm(embedding)

I used this code with just one GPU, and this works.
But this time I need to use more gpus, and I have to modify that tensor which was manually located device=self.device, to something “which is dynamically located each gpu at DataParallel”.
… But it is hard to me, code below doesn’t works. The tensor just located in cpu, even with .cuda()
How can I solve this issue? I am almost struggling with this issue all day…
class Embedding(nn.Module):
    def __init__(self, maxlen, d_model, n_segments, vocab_size, device='cuda'):
        super(Embedding, self).__init__()
        self.device = torch.device('cuda')
        self.tok_embed = nn.Embedding(vocab_size, d_model)  # token embedding
        self.pos_embed = nn.Embedding(maxlen, d_model).cpu()  # position embedding
        self.seg_embed = nn.Embedding(n_segments, d_model)  # segment(token type) embedding
        self.norm = nn.LayerNorm(d_model)
        self.pos = torch.arange(513, dtype=torch.long, requires_grad=False).unsqueeze(0)

    def forward(self, x, seg):
        # seq_len = x.size(1)
        print(f'pos device: {self.pos.device}') # printed by "cpu"
        pos = self.pos.expand_as(x)  # (seq_len,) -> (batch_size, seq_len)
        cuda_pos = self.pos_embed(pos).cuda()
        print(f'pos device: {cuda_pos.device}, x device : {x.device}, seg device: {seg.device}') # printed by "cpu" for pos
        embedding = self.tok_embed(x) + cuda_pos + self.seg_embed(seg)
        return self.norm(embedding)