How to apply model to multi-device with pytorch1.0?

I am reading the documentation of nn.DataParallel:

  • device_ids ( list of python:int or torch.device ) – CUDA devices (default: all devices)

And I can only set single device with code like device = torch.device("cuda:0").

How to set multi-device to use the default: all devices?

If you set device = None it uses all (None is the option by default in fact)

If I want to only use 4 gpus in a server with 8 gpus, how should I do that?

it’s better you to use the environment variable
CUDA_VISIBLE_DEVICES = id1,id2 and so on
it’s a cuda flag which will make any program to detect only those gpus you choose.

CUDA_VISIBLE_DEVICES=0,2,3,7 python your_script.py

The input tensor should also be put on gpu devices, I can do this with input_tensor.to(device) if there is only single gpu device. How should I do that with multi-gpu?

The input tensor should also be put on gpu devices, I can do this with input_tensor.to(device) if there is only single gpu device. How should I do that with multi-gpu?

If you are using dataparallel you don’t really need to move the tensor into devices, it’s done automatically once you feed the forward function.