From fa7a4f44b5c75410ebf5c0472c0e99cc1688a400 Mon Sep 17 00:00:00 2001
From: Maxim Vafin <maxim.vafin@intel.com>
Date: Thu, 30 Nov 2023 18:48:40 +0100
Subject: [PATCH] [PT FE] Update torch to 2.1 in GHA tests (#21391)

* [PT FE] Update torch to 2.1 in GHA tests

* Adjust requirements for natten

* Add upper bound

* Use sys.executable
---
 tests/constraints.txt                         |  2 +-
 .../torch_tests/requirements.txt              | 24 ++++++++++---------
 .../torch_tests/requirements_secondary.txt    |  4 +---
 .../torch_tests/test_detectron2.py            | 10 +++++++-
 4 files changed, 24 insertions(+), 16 deletions(-)

diff --git a/tests/constraints.txt b/tests/constraints.txt
index 79b06fe56e4..ec8bb52521c 100644
--- a/tests/constraints.txt
+++ b/tests/constraints.txt
@@ -23,4 +23,4 @@ pytest-html==4.1.1
 pytest-timeout==2.1.0
 jax<=0.4.14
 jaxlib<=0.4.14
-torch>=1.13
+torch>=1.13,<2.2
diff --git a/tests/model_hub_tests/torch_tests/requirements.txt b/tests/model_hub_tests/torch_tests/requirements.txt
index 470dd5073be..5954a485759 100644
--- a/tests/model_hub_tests/torch_tests/requirements.txt
+++ b/tests/model_hub_tests/torch_tests/requirements.txt
@@ -1,19 +1,21 @@
 -c ../../constraints.txt
+--extra-index-url https://download.pytorch.org/whl/cpu
+av
+basicsr
+datasets
+facexlib
 numpy
+pandas
+protobuf
+pyctcdecode
 pytest
 pytest-html
-torch<2.1
-torchvision
-av
-transformers
-timm
 sacremoses
 sentencepiece
-datasets
-pyctcdecode
-protobuf
 soundfile
-pandas
 super-image
-basicsr
-facexlib
+timm
+torch
+torchvision
+transformers
+wheel
diff --git a/tests/model_hub_tests/torch_tests/requirements_secondary.txt b/tests/model_hub_tests/torch_tests/requirements_secondary.txt
index 803c3c4b839..634f481dfc9 100644
--- a/tests/model_hub_tests/torch_tests/requirements_secondary.txt
+++ b/tests/model_hub_tests/torch_tests/requirements_secondary.txt
@@ -1,5 +1,3 @@
 -c ../../constraints.txt
 # This file contains requirements dependednt from modules in requirements.txt
-# get immutable commit to be able to use pip cache
-git+https://github.com/facebookresearch/detectron2.git@1a4df4d954089534daebbd7a0a45f08ece87056e
-natten
+natten -f https://shi-labs.com/natten/wheels/cpu/torch2.0.0/index.html 
diff --git a/tests/model_hub_tests/torch_tests/test_detectron2.py b/tests/model_hub_tests/torch_tests/test_detectron2.py
index 219c0ef5f1a..5c2ea9a1d9e 100644
--- a/tests/model_hub_tests/torch_tests/test_detectron2.py
+++ b/tests/model_hub_tests/torch_tests/test_detectron2.py
@@ -2,7 +2,8 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import os
-
+import sys
+import subprocess
 import pytest
 import torch
 from models_hub_common.utils import get_models_list, compare_two_tensors
@@ -19,6 +20,9 @@ class TestDetectron2ConvertModel(TestTorchConvertModel):
         self.image = Image.open(requests.get(url, stream=True).raw)
         self.image = self.image.resize([640, 480])
 
+        subprocess.run([sys.executable, "-m", "pip", "install",
+                       "git+https://github.com/facebookresearch/detectron2.git@017abbfa5f2c2a2afa045200c2af9ccf2fc6227f"])
+
     def load_model(self, model_name, model_link):
         from detectron2 import model_zoo, export
         from detectron2.modeling import build_model, PanopticFPN
@@ -82,6 +86,10 @@ class TestDetectron2ConvertModel(TestTorchConvertModel):
             is_ok = compare_two_tensors(cur_ov_res[:l], cur_fw_res[:l], fw_eps)
         assert is_ok, "Accuracy validation failed"
 
+    def teardown_class(self):
+        subprocess.run(
+            [sys.executable, "-m", "pip", "uninstall", "-y", "detectron2"])
+
     @pytest.mark.parametrize("name,type,mark,reason",
                              get_models_list(os.path.join(os.path.dirname(__file__), "detectron2_precommit")))
     @pytest.mark.precommit