The new .to(device) doesn't work with nn.DataParallel

autograd
1

the old code:

model=torch.nn.DataParallel(model,device_ids=range(torch.cuda.device_count()))
...
  for input, target in train_loader:
      input, target = Variable(input), Variable(target)
      input, target= input.cuda(), target.cuda()

this code works good
new code with .to(device)

model=torch.nn.DataParallel(model,device_ids=range(torch.cuda.device_count()))
...
  for input, target in train_loader:
       input, target = input.to(torch.cuda.current_device()), target.to(torch.cuda.current_device())

the input and target is only on GPU:0

1 Like
2

Or which should I pass to the device parameter when use .to(device) under dataparallel