Update sparseml to support torch 2.0 (#1618)

* bump up torch to 2.0, torchvision for compatability * bump torchvision * bump up torchaudio; test to make sure other package updates are compatible * remove comment * update patch version to support future torch 2.0.x versions
neuralmagic · Jun 13, 2023 · c3612a1 · c3612a1
1 parent ce0aaea
commit c3612a1
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/setup.py b/setup.py
@@ -62,17 +62,17 @@
 _deepsparse_ent_deps = [f"deepsparse-ent~={version_nm_deps}"]
 
 _onnxruntime_deps = ["onnxruntime>=1.0.0"]
-supported_torch_version = "torch>=1.7.0,<1.14"
+supported_torch_version = "torch>=1.7.0,<=2.0"
 _pytorch_deps = [
     supported_torch_version,
     "gputils",
 ]
 _pytorch_all_deps = _pytorch_deps + [
-    "torchvision>=0.3.0,<0.15",
-    "torchaudio<=0.13",
+    "torchvision>=0.3.0,<=0.15.1",
+    "torchaudio<=2.0.1",
 ]
 _pytorch_vision_deps = _pytorch_deps + [
-    "torchvision>=0.3.0,<0.15",
+    "torchvision>=0.3.0,<=0.15.1",
     "opencv-python<=4.6.0.66",
 ]
 _transformers_deps = _pytorch_deps + [

diff --git a/src/sparseml/pytorch/base.py b/src/sparseml/pytorch/base.py
@@ -49,7 +49,7 @@
 
 
 _TORCH_MIN_VERSION = "1.0.0"
-_TORCH_MAX_VERSION = "1.13.100"  # set bug to 100 to support all future 1.9.X versions
+_TORCH_MAX_VERSION = "2.0.100"
 
 
 def check_torch_install(