From 641408774ffd07afd4bfa0494022c3a544ae60b6 Mon Sep 17 00:00:00 2001
From: Benjamin <ben@neuralmagic.com>
Date: Wed, 31 Mar 2021 10:18:29 -0400
Subject: [PATCH] bump torch max supported version to 1.7

---
 README.md                                     | 2 +-
 integrations/ultralytics/README.md            | 4 ++--
 integrations/ultralytics/deepsparse/SERVER.md | 2 +-
 setup.py                                      | 4 ++--
 4 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index dbd6a8f10ef..85ebdd06004 100644
--- a/README.md
+++ b/README.md
@@ -312,7 +312,7 @@ clone the repository and install any additional dependencies as required.
 
 The currently supported framework versions are:
 
-- PyTorch supported versions: `>= 1.1.0, < 1.7.0`
+- PyTorch supported versions: `>= 1.1.0, < 1.8.0`
 - Keras supported versions: `2.3.0-tf` (through the TensorFlow `2.2` package; as of Feb 1st, 2021, `keras2onnx` has
 not been tested for TensorFlow >= `2.3`). 
 - TensorFlow V1 supported versions: >= `1.8.0` (TensorFlow >= `2.X` is not currently supported)
diff --git a/integrations/ultralytics/README.md b/integrations/ultralytics/README.md
index ade7a4ae19a..0b944ca609f 100644
--- a/integrations/ultralytics/README.md
+++ b/integrations/ultralytics/README.md
@@ -44,8 +44,8 @@ cp ../sparseml/integrations/ultralytics/*.py .
 cp ../sparseml/integrations/ultralytics/deepsparse/*.py .
 
 # install dependencies
+pip install sparseml[torchvision] deepsparse
 pip install -r requirements.txt
-pip install sparseml deepsparse
 ```
 
 
@@ -119,7 +119,7 @@ cp sparseml/integrations/ultralytics/deepsparse/*.py yolov5
 cd yolov5
 
 # install deepsparse and server dependencies
-pip install deepsparse sparseml flask flask-cors
+pip install deepsparse sparseml[torchvision] flask flask-cors
 ```
 
 Note: on new Ubuntu systems, to install `cv2` running `sudo apt-get update && apt-get install -y python3-opencv`
diff --git a/integrations/ultralytics/deepsparse/SERVER.md b/integrations/ultralytics/deepsparse/SERVER.md
index 04a7e1c1418..d740b33dcb4 100644
--- a/integrations/ultralytics/deepsparse/SERVER.md
+++ b/integrations/ultralytics/deepsparse/SERVER.md
@@ -41,8 +41,8 @@ cp sparseml/integrations/ultralytics/deepsparse/*.py yolov5
 cd yolov5
 
 # install dependencies
+pip install deepsparse sparseml[torchvision] flask flask-cors
 pip install -r requirements.txt
-pip install deepsparse sparseml flask flask-cors
 ```
 
 ## Execution
diff --git a/setup.py b/setup.py
index 6fe91eecac5..10a56135e13 100644
--- a/setup.py
+++ b/setup.py
@@ -55,8 +55,8 @@
 _deepsparse_deps = [
     f"{'deepsparse-nightly' if _NIGHTLY else 'deepsparse'}~={_VERSION_MAJOR_MINOR}"
 ]
-_pytorch_deps = ["torch>=1.1.0", "tensorboard>=1.0", "tensorboardX>=1.0"]
-_pytorch_vision_deps = _pytorch_deps + ["torchvision>=0.3.0"]
+_pytorch_deps = ["torch>=1.1.0,<1.8", "tensorboard>=1.0", "tensorboardX>=1.0"]
+_pytorch_vision_deps = _pytorch_deps + ["torchvision>=0.3.0,<0.9"]
 _tensorflow_v1_deps = ["tensorflow<2.0.0", "tensorboard<2.0.0", "tf2onnx>=1.0.0,<1.6"]
 _tensorflow_v1_gpu_deps = [
     "tensorflow-gpu<2.0.0",