Merge branch 'main' into feature/damian/sparsetensors

neuralmagic · Apr 22, 2024 · dfa41fb · dfa41fb
2 parents 5d9c7dd + 1abc30c
commit dfa41fb
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 5 deletions.
diff --git a/docs/source/installation.md b/docs/source/installation.md
@@ -63,7 +63,5 @@ TensorFlow V1 with GPU operations enabled:
 pip install sparseml[tf_v1_gpu]
 ```
 
-Depending on your device and CUDA version, you may need to install additional dependencies for using TensorFlow V1 with GPU operations. You can find these steps [here.](https://www.tensorflow.org/install/gpu#older_versions_of_tensorflow)
-
 Note, TensorFlow V1 is no longer being built for newer operating systems such as Ubuntu 20.04. 
 Therefore, SparseML with TensorFlow V1 is unsupported on these operating systems as well.
diff --git a/setup.py b/setup.py
@@ -68,7 +68,7 @@
 
 _onnxruntime_deps = ["onnxruntime>=1.0.0"]
 _clip_deps = ["open_clip_torch==2.20.0"]
-supported_torch_version = "torch>=1.7.0,<2.2"
+supported_torch_version = "torch>=1.7.0,<2.3"
 _pytorch_deps = [
     supported_torch_version,
     "gputils",

diff --git a/src/sparseml/pytorch/base.py b/src/sparseml/pytorch/base.py
@@ -49,7 +49,8 @@
 
 
 _TORCH_MIN_VERSION = "1.0.0"
-_TORCH_MAX_VERSION = os.environ.get("MAX_TORCH", "2.1.10")
+# set max to 2.2.99 to account for bugfix versions with 2.2
+_TORCH_MAX_VERSION = os.environ.get("MAX_TORCH", "2.2.99")
 
 
 def check_torch_install(

diff --git a/tests/sparseml/transformers/sparsification/modification/conftest.py b/tests/sparseml/transformers/sparsification/modification/conftest.py
@@ -67,7 +67,7 @@ def distilbert_model():
 
 @pytest.fixture
 def mistral_model():
-    config = AutoConfig.from_pretrained("mistralai/Mistral-7B-v0.1")
+    config = AutoConfig.from_pretrained("NousResearch/Hermes-2-Pro-Mistral-7B")
     with init_empty_weights():
         model = AutoModel.from_config(config)
     return model