@ptrblck Thanks a lot for this. I tried with your code and it works well only on the first step, and then the error occurs again. Here is the full error message:
torch.Size([5, 600, 111]) cuda:0
torch.Size([5, 600, 111]) cuda:1
torch.Size([5, 600, 111]) cuda:2
torch.Size([5, 600, 111]) cuda:3
Stage 1, epoch: 0,step: 0, loss: 0.5405522584915161
torch.Size([2, 600, 111]) cuda:0
torch.Size([2, 600, 111]) cuda:1
torch.Size([1, 600, 111]) cuda:2
Traceback (most recent call last):
File “e2e_learning_stage1.py”, line 349, in 
output = contact_net(PE_batch,seq_embedding_batch, state_pad)
File “/g/data/ik06/jiajia/python3packages/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 727, in _call_impl
result = self.forward(*input, **kwargs)
File “/g/data/ik06/jiajia/python3packages/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py”, line 161, in forward
outputs = self.parallel_apply(replicas, inputs, kwargs)
File “/g/data/ik06/jiajia/python3packages/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py”, line 171, in parallel_apply
return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
File “/g/data/ik06/jiajia/python3packages/lib/python3.8/site-packages/torch/nn/parallel/parallel_apply.py”, line 86, in parallel_apply
output.reraise()
File “/g/data/ik06/jiajia/python3packages/lib/python3.8/site-packages/torch/_utils.py”, line 428, in reraise
raise self.exc_type(msg)
RuntimeError: Caught RuntimeError in replica 0 on device 0.
Original Traceback (most recent call last):
File “/g/data/ik06/jiajia/python3packages/lib/python3.8/site-packages/torch/nn/parallel/parallel_apply.py”, line 61, in _worker
output = module(input, **kwargs)
File “/g/data/ik06/jiajia/python3packages/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 727, in _call_impl
result = self.forward(input, **kwargs)
File “/g/data/ik06/jiajia/python3packages/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py”, line 161, in forward
outputs = self.parallel_apply(replicas, inputs, kwargs)
File “/g/data/ik06/jiajia/python3packages/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py”, line 171, in parallel_apply
return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
File “/g/data/ik06/jiajia/python3packages/lib/python3.8/site-packages/torch/nn/parallel/parallel_apply.py”, line 86, in parallel_apply
output.reraise()
File “/g/data/ik06/jiajia/python3packages/lib/python3.8/site-packages/torch/_utils.py”, line 428, in reraise
raise self.exc_type(msg)
RuntimeError: Caught RuntimeError in replica 1 on device 1.
Original Traceback (most recent call last):
File “/g/data/ik06/jiajia/python3packages/lib/python3.8/site-packages/torch/nn/parallel/parallel_apply.py”, line 61, in _worker
output = module(input, **kwargs)
File “/g/data/ik06/jiajia/python3packages/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 727, in _call_impl
result = self.forward(input, **kwargs)
File “/g/data/ik06/jiajia/e2efold_master/e2efold/models.py”, line 252, in forward
position_embeds = self.PE_net(pe.view(-1, 111)).view(-1, self.L, self.d) # NL111 -> NLd
File “/g/data/ik06/jiajia/python3packages/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 727, in _call_impl
result = self.forward(*input, **kwargs)
File “/g/data/ik06/jiajia/python3packages/lib/python3.8/site-packages/torch/nn/modules/container.py”, line 117, in forward
input = module(input)
File “/g/data/ik06/jiajia/python3packages/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 727, in _call_impl
result = self.forward(*input, **kwargs)
File “/g/data/ik06/jiajia/python3packages/lib/python3.8/site-packages/torch/nn/modules/linear.py”, line 93, in forward
return F.linear(input, self.weight, self.bias)
File “/g/data/ik06/jiajia/python3packages/lib/python3.8/site-packages/torch/nn/functional.py”, line 1690, in linear
ret = torch.addmm(bias, input, weight.t())
RuntimeError: Expected tensor for ‘out’ to have the same device as tensor for argument #2 ‘mat1’; but device 0 does not equal 1 (while check
ing arguments for addmm)
I am using batch size 20, and 4 GPUs. I will also attach my training code here if helpful:
for epoch in range(epoches_first):
    for contacts, seq_embeddings, matrix_reps, seq_lens in train_generator:
        contact_net.train()
        contacts_batch = torch.Tensor(contacts.float()).cuda()
        seq_embedding_batch = torch.Tensor(seq_embeddings.float()).cuda()
        matrix_reps_batch = torch.unsqueeze(torch.Tensor(matrix_reps.float()).cuda(), -1)
        state_pad = torch.zeros([matrix_reps_batch.shape[0],seq_len, seq_len]).cuda()
        PE_batch = get_pe(seq_lens, seq_len).float().cuda()
        contact_masks = torch.Tensor(contact_map_masks(seq_lens, seq_len)).cuda()

        contact_net = torch.nn.DataParallel(contact_net)
        output = contact_net(PE_batch,seq_embedding_batch, state_pad)
        # Compute loss
        loss_u = criterion_bce_weighted(output*contact_masks, contacts_batch)
        # print(steps_done)
        if steps_done % OUT_STEP ==0:
            print('Stage 1, epoch: {},step: {}, loss: {}'.format(
                epoch, steps_done, loss_u))

        # Optimize the model
        u_optimizer.zero_grad()
        loss_u.backward()
        u_optimizer.step()
        steps_done=steps_done+1

Just to add the torch_seed() method I used to make sure it’s not causing trouble:
def seed_torch(seed=0):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed) # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True