Softmax dimension out of range with parallel training

HI, I try to use F.softmax(x, dim=1) with x is a 2-dim tensor on GPU, the code works well while I use a single GPU, but I get an error message ‘IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)’ after I use multi-GPU devices to train the network. could you help me to fix the issue?