Multi gpu train

[YaoooLiang]

HI,
when I train with multi gpu, I got that error:

Traceback (most recent call last):
File "train.py", line 231, in
main()
File "train.py", line 189, in main
soft_preds = torch.cat(soft_parsing, dim=0)
RuntimeError: All input tensors must be on the same device. Received cuda:0 and cuda:1

[pshyg]

I had same the problem.
I solved the problem in the following way.

In train.py change line 186, 187
186: soft_parsing.append(soft_pred[0][-1])
187: soft_edge.append(soft_pred[1][-1])

to

186: soft_parsing.append(soft_pred[0][-1].to(torch.device('cuda:0')))
187: soft_edge.append(soft_pred[1][-1].to(torch.device('cuda:0')))