A better way to make pytorch code agnostic to running on a CPU or GPU?
The Migration guide recommends the following to make code CPU/GPU agnostic:

&gt; # 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?

How about

if torch.cuda.device_count() &gt; 1:
    model = torch.nn.DataParallel(model)
model = model.to(device)


?

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