Pytorch Imagenet Github Example - Training input variable not cuda?

In the github example located here: https://github.com/pytorch/examples/blob/master/imagenet/main.py#L163

Why is the input variable not moved to cuda?

Thank you!

If your input has to use 2 GPUs, it’s more efficient to send the first half of the input to GPU1 and second half to GPU2, rather than sending the entire input to GPU1, and then sending half of it from GPU1 to GPU2.

That’s the reason that the input is not transferred to the GPU at the code location that you pointed out.

That example is multi-GPU ready – the model is wrapped in an nn.DataParallel, and nn.DataParallel broadcasts the input living on the CPU efficiently to the number of GPUs being used.

5 Likes