From e8c6be1f40fe67f8584a492e60fa33b359347f0a Mon Sep 17 00:00:00 2001
From: Matthias Gehre <matthias.gehre@amd.com>
Date: Mon, 10 Jun 2024 22:07:24 +0200
Subject: [PATCH] Update torch stable

---
 projects/pt1/e2e_testing/xfail_sets.py                    | 7 -------
 projects/pt1/python/torch_mlir/dynamo.py                  | 4 +---
 .../pt1/python/torch_mlir_e2e_test/test_suite/__init__.py | 8 --------
 stable-requirements.txt                                   | 4 ++--
 4 files changed, 3 insertions(+), 20 deletions(-)

diff --git a/projects/pt1/e2e_testing/xfail_sets.py b/projects/pt1/e2e_testing/xfail_sets.py
index a8e4649a96b8..2d24012f2b6c 100644
--- a/projects/pt1/e2e_testing/xfail_sets.py
+++ b/projects/pt1/e2e_testing/xfail_sets.py
@@ -13,8 +13,6 @@
 from torch_mlir_e2e_test.test_suite import COMMON_TORCH_MLIR_LOWERING_XFAILS
 from torch_mlir._version import torch_version_for_comparison, version
 
-print(f"TORCH_VERSION_FOR_COMPARISON =", torch_version_for_comparison())
-
 LINALG_XFAIL_SET = COMMON_TORCH_MLIR_LOWERING_XFAILS | {
     "Conv1dNoPaddingGroupModule_basic",
     "RepeatInterleaveStaticModule_basic",
@@ -352,11 +350,6 @@
     "InterpolateDynamicModule_scales_recompute_bilinear",
 }
 
-if torch_version_for_comparison() <= version.parse("2.2.0"):
-    TORCHDYNAMO_XFAIL_SET |= {
-        'OneHotModule_basic',
-    }
-
 TORCHDYNAMO_CRASHING_SET = {
     # No upstream decompositions.
     # %6:4 = torch.operator "aten._embedding_bag_forward_only"(%1, %3, %5, %false, %int0, %false, %none, %false, %int-1) : (!torch.tensor<*,f32>, !torch.tensor<*,si64>, !torch.tensor<*,si64>, !torch.bool, !torch.int, !torch.bool, !torch.none, !torch.bool, !torch.int) -> (!torch.tensor, !torch.tensor, !torch.tensor, !torch.tensor)
diff --git a/projects/pt1/python/torch_mlir/dynamo.py b/projects/pt1/python/torch_mlir/dynamo.py
index 7fc887d56bc4..7622ec013659 100644
--- a/projects/pt1/python/torch_mlir/dynamo.py
+++ b/projects/pt1/python/torch_mlir/dynamo.py
@@ -50,6 +50,7 @@ def _get_decomposition_table():
         # (the upstream decomposition we use here does), even though we have
         # support for aten.native_batch_norm_backward.
         aten._native_batch_norm_legit_functional,
+        aten._native_batch_norm_legit_no_training,
         aten.native_group_norm,
         aten.split.Tensor,
         aten.split_with_sizes,
@@ -67,9 +68,6 @@ def _get_decomposition_table():
         aten.cumsum,
         aten.index_select,
     ]
-    # TODO: enable test once 2.1.0 is stable
-    if torch_version_for_comparison() >= version.parse("2.1.0.dev"):
-        decomp_list += [aten._native_batch_norm_legit_no_training]
     return get_decompositions(decomp_list)
 
 
diff --git a/projects/pt1/python/torch_mlir_e2e_test/test_suite/__init__.py b/projects/pt1/python/torch_mlir_e2e_test/test_suite/__init__.py
index 6f492a1eff5c..c03fd95505a6 100644
--- a/projects/pt1/python/torch_mlir_e2e_test/test_suite/__init__.py
+++ b/projects/pt1/python/torch_mlir_e2e_test/test_suite/__init__.py
@@ -18,14 +18,6 @@
     "ElementwiseToDtypeI64ToUI8Module_basic",
 }
 
-# TODO: Delete once torch 2.1.0 is released
-if torch_version_for_comparison() < version.parse("2.1.0.dev"):
-    COMMON_TORCH_MLIR_LOWERING_XFAILS.update({
-        "ScaledDotProductAttentionDifferentModule_basic",
-        "ScaledDotProductAttentionSameModule_basic"
-    })
-
-
 def register_all_tests():
     """Registers all the built-in E2E tests that Torch-MLIR provides."""
     # Side-effecting import statements.
diff --git a/stable-requirements.txt b/stable-requirements.txt
index 1641e0540671..27d0c30d7a91 100644
--- a/stable-requirements.txt
+++ b/stable-requirements.txt
@@ -1,3 +1,3 @@
 --index-url https://download.pytorch.org/whl/cpu
-torch==2.1.2+cpu
-torchvision==0.16.2+cpu
+torch==2.3.1+cpu
+torchvision==0.18.1+cpu