完整报错信息
ValueError: DistributedDataParallel device_ids and output_device arguments only work with single-device/multiple-device GPU modules or CPU modules, but got device_ids [1], output_device 1, and module parameters {device(type=‘cpu’)}.
原因
nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)
? 在使用“DistributedDataParallel”时,model没有先放在GPU上,而是默认的CPU。只需要如下修改既可。
nn.parallel.DistributedDataParallel(model.cuda(), device_ids=[local_rank], output_device=local_rank)