nn.DataParallel
splits the data along the batch dimension so that each specified GPU will get a chunk of the batch. If you just call .cuda()
(or the equivalent .to()
call), you will push the tensor or parameters onto the specified single device.
2 Likes