已合并
[众智][PyTorch]3DUNet模型添加DDP参数find_unused_parameters=True #4707
创建于 2023年5月17日
[众智][PyTorch]3DUNet模型添加DDP参数find_unused_parameters=True #4707
已合并
从refs/pull/4707/head合入到master
共 1 个文件变更+1-1
| @@ -61,7 +61,7 @@ def main(): | |||
| 61 | model, optimizer = amp.initialize(model, optimizer, opt_level="O1", loss_scale=None, combine_grad=True) | 61 | model, optimizer = amp.initialize(model, optimizer, opt_level="O1", loss_scale=None, combine_grad=True) |
| 62 | 62 | ||
| 63 | if args.world_size > 1: | 63 | if args.world_size > 1: |
| 64 | model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.rank]) | 64 | model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.rank], find_unused_parameters=True) |
| 65 | 65 | ||
| 66 | 66 | ||
| 67 | trainer = train.Trainer(args, model, criterion, optimizer, train_data_loader=training_generator, | 67 | trainer = train.Trainer(args, model, criterion, optimizer, train_data_loader=training_generator, |