torchtnt.utils.prepare_module.prepare_ddp

torchtnt.utils.prepare_module.prepare_ddp(module: Module, device: device, strategy: Optional[DDPStrategy] = None) → DistributedDataParallel

Utility to move a module to device and wrap in DistributedDataParallel.

Parameters:

• module – module to be wrapped in DDP
• device – device to which module will be moved
• strategy – an instance of DDPStrategy which defines the settings of DDP APIs

Examples::

strategy = DDPStrategy(find_unused_parameters=True, gradient_as_bucket_view=True) module = nn.Linear(1, 1) device = torch.device(“cuda”) ddp_module = prepare_ddp(module, device, strategy)