When running on GPU, and constructing the list of inputs, they might not end
up on the right device (i.e. the input is on gpu instead of cpu or the other
way around). The simple fix is just to filter the arguments.