Make py3nvml optional

rasbt · Jul 1, 2023 · 5fe85f2 · 5fe85f2
1 parent b9b51dd
commit 5fe85f2
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 3 deletions.
diff --git a/requirements.txt b/requirements.txt
@@ -1,3 +1,2 @@
 ipython >= 6.0
 importlib-metadata >= 1.4
-py3nvml >= 0.2
diff --git a/setup.py b/setup.py
@@ -32,6 +32,7 @@
     url="https://github.com/rasbt/watermark",
     packages=find_packages(exclude=[]),
     install_requires=install_reqs,
+    extras_require={'gpu': ['py3nvml>=0.2']},
     long_description=dedent(
         """\
         An IPython magic extension for printing date and time stamps, version

diff --git a/watermark/watermark.py b/watermark/watermark.py
@@ -19,7 +19,11 @@
 from multiprocessing import cpu_count
 from socket import gethostname
 import platform
-from py3nvml import py3nvml
+
+try:
+    from py3nvml import py3nvml
+except ImportError:
+    py3nvml = None
 
 try:
     import importlib.metadata as importlib_metadata
@@ -350,4 +354,4 @@ def _get_gpu_info():
                 "to be installed on this machine."}
     except:
         return {"GPU Info": "GPU information is not "
-                "available for this machine."}
+                "available for this machine."}