feat: add the option to use a seingel gpu

PHT-Medic · Nov 14, 2023 · c62dde9 · c62dde9
1 parent da7ba37
commit c62dde9
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/airflow/dags/DAG_PHT_run_train.py b/airflow/dags/DAG_PHT_run_train.py
@@ -273,6 +273,10 @@ def execute_container(train_state):
                     raise ValueError(
                         f"Invalid gpu configuration: {gpu_config}. Must be a list of integers or 'all'"
                     )
+            elif isinstance(gpu_config, int):
+                device_request = docker.types.DeviceRequest(
+                    device_ids=[str(gpu_config)], capabilities=[["gpu"]]
+                )
             else:
                 raise ValueError(
                     f"Invalid gpu configuration: {gpu_config}. Must be a list of integers or 'all'"