From 81aebeeaf64a0cb490a6d54e35bd63bc233b2f92 Mon Sep 17 00:00:00 2001
From: David Jurado <david.jurado@factored.ai>
Date: Wed, 15 Nov 2023 15:02:04 -0500
Subject: [PATCH] Add CUDA_VISIBLE_DEVICES enviroment variable when using the
 --gpus flag

---
 mlcube/mlcube/parser.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/mlcube/mlcube/parser.py b/mlcube/mlcube/parser.py
index 17ee426..0ef794d 100644
--- 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[