Multi-GPU for training

I want to train my model on two gpus and I have set my model to:

model = torch.nn.DataParallel(model).cuda()

However, I notice that only gpu 0 is running and gpu 1 is not doing any work. Did I missed anything?


1 Like

Before asking any question related to code, I advise to look into documentation. In your case exactly: DataParallel

net = torch.nn.DataParallel(model, device_ids=[0, 1, 2])

1 Like

Thanks but it did not work. The gpus with index 1 and 2 did not run any python processes.

How do you check that above command is not working? I’m working with 2 GPU’s and it works very nice checked by nvidia-smi

You also want to make sure you set your environment variable: CUDA_VISIBLE_DEVICES=0,1,2 (or however many GPUs you are intending to use and have available)