Thanks for your reply! @colesbury
I tried what said. It works. However, Here comes another related problem. I also need to use the optimizer state saved in the previous single-gpu training rather than create a new optimizer.
optimizer.load_state_dict(checkpoint['optimizer'])

When I do this, the same error occurs. I guess something inside the saved optimizer is not consistent with the multi-gpus setting. Any suggestion to solve this? Thanks.