Bug fix for GPUs

ludwig-ai · Jun 6, 2022 · 2589297 · 2589297
1 parent 3999703
commit 2589297
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ludwig/hyperopt/execution.py b/ludwig/hyperopt/execution.py
@@ -476,7 +476,7 @@ def on_trainer_train_teardown(self, trainer, progress_tracker, save_path, is_coo
                 "use_gpu": use_gpu,
                 "resources_per_worker": {
                     "CPU": num_cpus,
-                    "GPU": 1,
+                    "GPU": 1 if use_gpu else 0,
                 },
             }
             hyperopt_dict["backend"].set_distributed_kwargs(**hvd_kwargs)