Update multi_node_cluster_auto_slurm.py

2019-09-09 10:55:47 -04:00 · 2019-09-09 10:55:47 -04:00 · ac0111c196
parent cbc619afa1
commit ac0111c196
1 changed files with 1 additions and 2 deletions
--- a/examples/new_project_templates/multi_node_examples/multi_node_cluster_auto_slurm.py
+++ b/examples/new_project_templates/multi_node_examples/multi_node_cluster_auto_slurm.py
@ -76,12 +76,11 @@ def main(hparams, cluster):
    # ------------------------
    # 4 INIT TRAINER
    # ------------------------
-    gpus = list(range(0, hparams.per_experiment_nb_gpus))
    trainer = Trainer(
        experiment=exp,
        checkpoint_callback=checkpoint,
        early_stop_callback=early_stop,
-        gpus=gpus,
+        gpus=hparams.per_experiment_nb_gpus,
        nb_gpu_nodes=hyperparams.nb_gpu_nodes
    )