From 3e1396eb18b33ccb387408bf436c672f50466b69 Mon Sep 17 00:00:00 2001
From: Daniel King <43149077+dakinggg@users.noreply.github.com>
Date: Mon, 17 Jun 2024 08:52:48 -0700
Subject: [PATCH] Add pynvml to mlflow dep group (#3404)

---
 setup.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/setup.py b/setup.py
index 2eba4e39e6..5beb4d5136 100644
--- a/setup.py
+++ b/setup.py
@@ -225,6 +225,7 @@ def package_files(prefix: str, directory: str, extension: str):
 extra_deps['mlflow'] = [
     'mlflow>=2.11.1,<3.0',
     'databricks-sdk==0.28.0',
+    'pynvml>=11.5.0,<12',
 ]
 
 extra_deps['pandas'] = ['pandas>=2.0.0,<3.0']