cybaj:

    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)
        pos.requires_grad = False
        print(f'pos tensor device: {pos.device}, shape: {pos.shape}')
        embedding = self.tok_embed(x) + self.pos_embed(pos) + self.seg_embed(seg)
        return self.norm(embedding)


pos tensor device: cuda:1, shape: torch.Size([64, 513])
pos tensor device: cuda:0, shape: torch.Size([64, 513])



Hey @cybaj, I am a little confused about the above code. IIUC, the self.device attribute on both replicas point to cuda:0, as replicate.py is not smart enough to  change that for you. In that case how did pos successfully located to different devices? Looks like expand_as would not automatically change device either:
>>> import torch
>>> x = torch.arange(10, device="cuda:0")
>>> y = torch.ones(10, 10).to(1)
>>> z = x.expand_as(y)
>>> z.device
device(type='cuda', index=0)
>>> y.device
device(type='cuda', index=1)
>>> x.device
device(type='cuda', index=0)

If you would like to get the correct device, can you read that from x.device? (input to forward should be scattered properly.)