diff --git a/examples/new_project_templates/single_gpu_node_ddp_template.py b/examples/new_project_templates/single_gpu_node_ddp_template.py index 95f69266..b5966a0e 100644 --- a/examples/new_project_templates/single_gpu_node_ddp_template.py +++ b/examples/new_project_templates/single_gpu_node_ddp_template.py @@ -71,7 +71,7 @@ def main(hparams): checkpoint_callback=checkpoint, early_stop_callback=early_stop, gpus=hparams.gpus, - distributed_backend='ddp' + distributed_backend=hparams.dist_backend ) # ------------------------ @@ -95,6 +95,8 @@ if __name__ == '__main__': parent_parser.add_argument('--gpus', type=str, default='-1', help='how many gpus to use in the node.' ' value -1 uses all the gpus on the node') + parent_parser.add_argument('--dist_backend', type=str, default='ddp', + help='When using multiple GPUs set Trainer(distributed_backend=dp) (or ddp)') parent_parser.add_argument('--test_tube_save_path', type=str, default=test_tube_dir, help='where to save logs') parent_parser.add_argument('--model_save_path', type=str, default=checkpoint_dir,