Add CUDA_VISIBLE_DEVICES enviroment variable when using the --gpus flag

mlcommons · Nov 15, 2023 · 81aebee · 81aebee
1 parent 755d388
commit 81aebee
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/mlcube/mlcube/parser.py b/mlcube/mlcube/parser.py
@@ -128,6 +128,9 @@ def parse_extra_arg(
             if parsed_args.get("gpus", None):
                 if platform == "docker":
                     runner_run_args["--gpus"] = parsed_args["gpus"]
+                    os.environ["CUDA_VISIBLE_DEVICES"] = parsed_args[
+                        "gpus"
+                    ]
                 else:
                     runner_run_args["--nv"] = ""
                     os.environ["SINGULARITYENV_CUDA_VISIBLE_DEVICES"] = parsed_args[