From 9651855ca431153ccf4c9439725d4e5a47601b5e Mon Sep 17 00:00:00 2001 From: Aayam Date: Fri, 28 Oct 2022 16:38:57 -0700 Subject: [PATCH] Add environment variable MUJOCO_PY_FORCE_GPU Currently I cannot find a way to show custom nvidia drive files for mujoco_py. hence I am suggesting to add this environment variable to force gpu as a failsafe measure to run in clusters. --- mujoco_py/builder.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/mujoco_py/builder.py b/mujoco_py/builder.py index 5280b2c3..dd672076 100644 --- a/mujoco_py/builder.py +++ b/mujoco_py/builder.py @@ -72,9 +72,12 @@ def load_cython_ext(mujoco_path): Builder = MacExtensionBuilder elif sys.platform == 'linux': _ensure_set_env_var("LD_LIBRARY_PATH", lib_path) + if os.getenv('MUJOCO_PY_FORCE_CPU') is None and get_nvidia_lib_dir() is not None: _ensure_set_env_var("LD_LIBRARY_PATH", get_nvidia_lib_dir()) Builder = LinuxGPUExtensionBuilder + elif os.getenv('MUJOCO_PY_FORCE_GPU') is not None: + Builder = LinuxGPUExtensionBuilder else: Builder = LinuxCPUExtensionBuilder elif sys.platform.startswith("win"):