From 35e80a836888eae6339b5ac81a49e489da2a12d9 Mon Sep 17 00:00:00 2001
From: Ben Greiner <code@bnavigator.de>
Date: Sat, 1 Jul 2023 18:47:45 +0200
Subject: [PATCH] Make py3nvml optional

---
 requirements.txt       | 1 -
 setup.py               | 1 +
 watermark/watermark.py | 8 ++++++--
 3 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index 002d861..55712cf 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,2 @@
 ipython >= 6.0
 importlib-metadata >= 1.4
-py3nvml >= 0.2
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 6d6428f..83a86c2 100644
--- 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
index be8d10b..b2542eb 100644
--- 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."}
\ No newline at end of file
+                "available for this machine."}