Add pynvml to mlflow dep group (#3404)

mosaicml · Jun 17, 2024 · 8119c14 · 8119c14
1 parent 4390dd6
commit 8119c14
Showing 1 changed file with 1 addition and 0 deletions.
diff --git 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']