def configure_ddp(model, rank): dist.init_process_group("nccl") model = DDP(model, device_ids=[rank], static_graph=True, gradient_as_bucket_view=True, bucket_cap_mb=2000) return model