configs/multi_gpu_train_diffusion.json

{
    "diffusion": {
        "_target_": "torch.nn.parallel.DistributedDataParallel",
        "module": "$@network_def.to(@device)",
        "device_ids": [
            "@device"
        ],
        "find_unused_parameters": true
    },
    "run": [
        "@load_autoencoder",
        "[email protected]()",
        "$print('scale factor:',@scale_factor)",
        "$@train#trainer.run()"
    ]
}