From 8e6781304c496da20f8614aad4c7630d2d79af15 Mon Sep 17 00:00:00 2001
From: Akshaya Agrawal <76990931+JBVAkshaya@users.noreply.github.com>
Date: Sat, 26 Jul 2025 13:21:43 -0700
Subject: [PATCH] Fix NVIDIA environment variables (#400)

(cherry picked from commit 812fea605cd2056ba20fac504f456457594595fb)
---
 .devcontainer/nvidia/devcontainer.json | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/.devcontainer/nvidia/devcontainer.json b/.devcontainer/nvidia/devcontainer.json
index 27f58aa6..38cf2f6f 100644
--- a/.devcontainer/nvidia/devcontainer.json
+++ b/.devcontainer/nvidia/devcontainer.json
@@ -14,6 +14,7 @@
     "--privileged",
     "--volume=/tmp/.X11-unix:/tmp/.X11-unix",
     "--volume=/mnt/wslg:/mnt/wslg",
+    "--runtime=nvidia",
     "--gpus=all"
   ],
   "containerEnv": {
@@ -22,14 +23,23 @@
     "XDG_RUNTIME_DIR": "${localEnv:XDG_RUNTIME_DIR}",
     "PULSE_SERVER": "${localEnv:PULSE_SERVER}",
     "LIBGL_ALWAYS_SOFTWARE": "1",
-    "QT_X11_NO_MITSHM": "1"
+    "QT_X11_NO_MITSHM": "1",
+    "NVIDIA_VISIBLE_DEVICES": "all",
+    "NVIDIA_DRIVER_CAPABILITIES": "all",
+    "__GLX_VENDOR_LIBRARY_NAME": "nvidia",
+    "__NV_PRIME_RENDER_OFFLOAD": "1",
+    "__VK_LAYER_NV_optimus": "NVIDIA_only"
   },
   "customizations": {
     "vscode": {
       "settings": {
         "python.defaultInterpreterPath": "/home/ubuntu/.venv/blue/bin/python",
-        "python.autoComplete.extraPaths": ["${workspaceFolder}/install/"],
-        "python.analysis.extraPaths": ["${workspaceFolder}/install/"]
+        "python.autoComplete.extraPaths": [
+          "${workspaceFolder}/install/"
+        ],
+        "python.analysis.extraPaths": [
+          "${workspaceFolder}/install/"
+        ]
       },
       "extensions": [
         "ms-azuretools.vscode-docker",