dmlcloud.wrap_ddp

dmlcloud.wrap_ddp(module, device, sync_bn=False, find_unused_parameters=False, verbose=True)

Wraps a module with DistributedDataParallel.

Parameters:

• module (nn.Module) – The module to wrap.

• device (torch.device) – The device to use.

• sync_bn (bool, optional) – If True, uses SyncBatchNorm. Default is False.

• find_unused_parameters (bool, optional) – If True, finds unused parameters. Default is False.

• verbose (bool, optional) – If True, prints information about the model. Default is True.

Returns:

The wrapped module.

Return type:

nn.Module