A better way to make pytorch code agnostic to running on a CPU or GPU?

Question

The Migration guide recommends the following to make code CPU/GPU agnostic:

> # at beginning of the script
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
...
# then whenever you get a new Tensor or Module
# this won't copy if they are already on the desired device
input = data.to(device)
model = MyModule(...).to(device)

I did this and ran my code on a CPU-only device, but my model crashed when fed an input array, as saying it was expecting a CPU tensor not a GPU one. Somehow my model was automatically converting the a CPU-input array to a GPU array. Finally I traced it down to this command in my code:

model = torch.nn.DataParallel(model).to(device)

Even though I convert the model to 'cpu', the nn.DataParallel overrides this. The best solution I came up with was a conditional:

if device.type=='cpu':
    model = model.to(device)
else:
    model = torch.nn.DataParallel(model).to(device)

This does not seem elegant. Is there a better way?

Answer 1

How about

if torch.cuda.device_count() > 1:
    model = torch.nn.DataParallel(model)
model = model.to(device)

?

You don't need DataParallel if you have only one GPU.