torchfoo.module.make_ddp

torchfoo.module.make_ddp(m)

Wrap a module with DDP and convert BatchNorm to SyncBatchNorm.

Returns the module unchanged if not in distributed mode.

Warning

SyncBatchNorm only works with GPU modules. On CPU, BatchNorm layers will not be converted and a warning is emitted.

Parameters:

m (Module) – the module to wrap

Returns:

The DDP-wrapped module, or the original module if not distributed.