[Python] DataParallel multi-gpu RuntimeError: chunk expects at least a 1-dimensional tensor

I am trying to run my model on multiple gpus using DataParallel by setting model = nn.DataParallel(model).cuda(), but everytime getting this error:

RuntimeError: chunk expects at least a 1-dimensional tensor (chunk at /pytorch/aten/src/ATen/native/TensorShape.cpp:184).
​

My code is correct. Does anyone know what's wrong?

I have tried setting device_ids=[0,1] parameter and also CUDA_VISIBLE_DEVICES on the terminal. Also tried different batch sizes.

Continue reading...