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