From 4575c404b97c6c5140a123529354a0c9c1bc3362 Mon Sep 17 00:00:00 2001
From: YouJiacheng <1503679330@qq.com>
Date: Tue, 19 Jul 2022 05:35:58 +0800
Subject: [PATCH] Support nvidia-docker v2

see
https://github.com/NVIDIA/nvidia-docker/issues/789
---
 mujoco_py/builder.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/mujoco_py/builder.py b/mujoco_py/builder.py
index 5280b2c3..eebfb0f9 100644
--- a/mujoco_py/builder.py
+++ b/mujoco_py/builder.py
@@ -28,9 +28,12 @@ def get_nvidia_lib_dir():
                                         stdout=subprocess.PIPE, stderr=subprocess.PIPE) == 0
     if not exists_nvidia_smi:
         return None
-    docker_path = '/usr/local/nvidia/lib64'
-    if exists(docker_path):
-        return docker_path
+    docker_path_v1 = '/usr/local/nvidia/lib64'
+    if exists(docker_path_v1):
+        return docker_path_v1
+    docker_path_v2 = '/usr/local/cuda/lib64'
+    if exists(docker_path_v2):
+        return docker_path_v2
 
     nvidia_path = '/usr/lib/nvidia'
     if exists(nvidia_path):