bump pydantic

setup.py

@@ -22,7 +22,7 @@ def _setup_packages() -> List:
     )
 
 def _setup_install_requires() -> List:
-    return ["torch>=1.7.0", "transformers<=4.40", "pydantic>=1.8.2,<2.0.0"]
+    return ["torch>=1.7.0", "transformers<=4.40", "pydantic<2.7"]
 
 def _setup_extras() -> Dict:
     return {"dev": ["black==22.12.0", "isort==5.8.0", "wheel>=0.36.2", "flake8>=3.8.3", "pytest>=6.0.0"]}

src/sparsetensors/README.md

@@ -65,7 +65,7 @@ is saved to the model's `config.json` file. **Note:** the model must have been
 initialized with SparseAutoModelForCausalLM.from_pretrained()
 
 ```python
-from sparseml.transformers.compression import BitmaskConfig
+from sparsetensors import BitmaskConfig
 
 output_dir = "/PATH/TO/SAVE/COMPRESSED_MODEL"
 sparsity_config = BitmaskConfig()