I follow your method and it works well when I only use one GPU.
However, when I try to run it on Dataparallel, it gives me the error:

RuntimeError: Expected tensor for argument #1 ‘input’ to have the same device as tensor for argument #2 ‘weight’; but device 1 does not equal 0 (while checking arguments for cudnn_convolution)

Any suggestion?
Thanks!