From cdac067bf17c87d286c2fef9d00994452e4a89dc Mon Sep 17 00:00:00 2001
From: Abdol A <u2271662@live.warwick.ac.uk>
Date: Thu, 21 Mar 2024 15:01:39 +0000
Subject: [PATCH 01/19] =?UTF-8?q?=E2=9A=A1=EF=B8=8F=20Add=20`torch.compile?=
 =?UTF-8?q?`=20to=20WSI=20registration?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 tiatoolbox/tools/registration/wsi_registration.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/tiatoolbox/tools/registration/wsi_registration.py b/tiatoolbox/tools/registration/wsi_registration.py
index cacfa0e38..6c8c9110b 100644
--- a/tiatoolbox/tools/registration/wsi_registration.py
+++ b/tiatoolbox/tools/registration/wsi_registration.py
@@ -16,7 +16,8 @@
 from skimage.util import img_as_float
 from torchvision.models import VGG16_Weights
 
-from tiatoolbox import logger
+from tiatoolbox import logger, rcParam
+from tiatoolbox.models.architecture.utils import compile_model
 from tiatoolbox.tools.patchextraction import PatchExtractor
 from tiatoolbox.utils.metrics import dice
 from tiatoolbox.utils.transforms import imresize
@@ -338,9 +339,11 @@ def __init__(self: torch.nn.Module) -> None:
         output_layers_id: list[str] = ["16", "23", "30"]
         output_layers_key: list[str] = ["block3_pool", "block4_pool", "block5_pool"]
         self.features: dict = dict.fromkeys(output_layers_key, None)
-        self.pretrained: torch.nn.Sequential = torchvision.models.vgg16(
-            weights=VGG16_Weights.IMAGENET1K_V1,
-        ).features
+        self.pretrained: torch.nn.Sequential = compile_model(
+            torchvision.models.vgg16(weights=VGG16_Weights.IMAGENET1K_V1),
+            mode=rcParam["torch_compile_mode"],
+            disable=not rcParam["enable_torch_compile"],
+            ).features
         self.f_hooks = [
             getattr(self.pretrained, layer).register_forward_hook(
                 self.forward_hook(output_layers_key[i]),

From b599b665bcb685931ff84b4c66501835f8d8681b Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 21 Mar 2024 15:06:20 +0000
Subject: [PATCH 02/19] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tiatoolbox/tools/registration/wsi_registration.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tiatoolbox/tools/registration/wsi_registration.py b/tiatoolbox/tools/registration/wsi_registration.py
index 6c8c9110b..d49f40620 100644
--- a/tiatoolbox/tools/registration/wsi_registration.py
+++ b/tiatoolbox/tools/registration/wsi_registration.py
@@ -343,7 +343,7 @@ def __init__(self: torch.nn.Module) -> None:
             torchvision.models.vgg16(weights=VGG16_Weights.IMAGENET1K_V1),
             mode=rcParam["torch_compile_mode"],
             disable=not rcParam["enable_torch_compile"],
-            ).features
+        ).features
         self.f_hooks = [
             getattr(self.pretrained, layer).register_forward_hook(
                 self.forward_hook(output_layers_key[i]),

From 56c4e121b879bed160ab2cea785891f03cc47d2f Mon Sep 17 00:00:00 2001
From: Abdol A <u2271662@live.warwick.ac.uk>
Date: Thu, 11 Apr 2024 10:50:10 +0100
Subject: [PATCH 03/19] =?UTF-8?q?=E2=9C=85=20Add=20DFBR=20with=20torch.com?=
 =?UTF-8?q?pile=20test?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 tests/models/test_patch_predictor.py |  3 +-
 tests/test_wsi_registration.py       | 68 ++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/tests/models/test_patch_predictor.py b/tests/models/test_patch_predictor.py
index 6c04c6b68..5b94ae189 100644
--- a/tests/models/test_patch_predictor.py
+++ b/tests/models/test_patch_predictor.py
@@ -1239,7 +1239,7 @@ def test_patch_predictor_torch_compile(
     sample_patch2: Path,
     tmp_path: Path,
 ) -> None:
-    """Test torch.compile functionality.
+    """Test PatchPredictor with with torch.compile functionality.
 
     Args:
         sample_patch1 (Path): Path to sample patch 1.
@@ -1250,7 +1250,6 @@ def test_patch_predictor_torch_compile(
     torch_compile_enabled = rcParam["enable_torch_compile"]
     torch._dynamo.reset()
     rcParam["enable_torch_compile"] = True
-    # Test torch.compile with default mode
     rcParam["torch_compile_mode"] = "default"
     _, compile_time = timed(
         test_patch_predictor_api,
diff --git a/tests/test_wsi_registration.py b/tests/test_wsi_registration.py
index 79abd3855..470cb281c 100644
--- a/tests/test_wsi_registration.py
+++ b/tests/test_wsi_registration.py
@@ -5,7 +5,10 @@
 import cv2
 import numpy as np
 import pytest
+import torch
 
+from tests.conftest import timed
+from tiatoolbox import logger, rcParam
 from tiatoolbox.tools.registration.wsi_registration import (
     AffineWSITransformer,
     DFBRegister,
@@ -576,3 +579,68 @@ def test_affine_wsi_transformer(sample_ome_tiff: Path) -> None:
         expected = cv2.rotate(expected, cv2.ROTATE_90_CLOCKWISE)
 
         assert np.sum(expected - output) == 0
+
+def test_dfbr_feature_extractor_torch_compile(dfbr_features: Path) -> None:
+    """Test DFBRFeatureExtractor with torch.compile functionality.
+
+    Args:
+        dfbr_features (Path): Path to the expected features.
+    """
+    torch_compile_enabled = rcParam["enable_torch_compile"]
+    torch._dynamo.reset()
+    rcParam["enable_torch_compile"] = True
+
+    def _extract_features() -> tuple:
+        dfbr = DFBRegister()
+        fixed_img = np.repeat(
+            np.expand_dims(
+                np.repeat(
+                    np.expand_dims(np.arange(0, 64, 1, dtype=np.uint8), axis=1),
+                    64,
+                    axis=1,
+                ),
+                axis=2,
+            ),
+            3,
+            axis=2,
+        )
+        output = dfbr.extract_features(fixed_img, fixed_img)
+        pool3_feat = output["block3_pool"][0, :].detach().numpy()
+        pool4_feat = output["block4_pool"][0, :].detach().numpy()
+        pool5_feat = output["block5_pool"][0, :].detach().numpy()
+
+        return pool3_feat, pool4_feat, pool5_feat
+
+    (pool3_feat, pool4_feat, pool5_feat), compile_time = timed(_extract_features)
+    _pool3_feat, _pool4_feat, _pool5_feat = np.load(
+        str(dfbr_features),
+        allow_pickle=True,
+    )
+    assert np.mean(np.abs(pool3_feat - _pool3_feat)) < 1.0e-4
+    assert np.mean(np.abs(pool4_feat - _pool4_feat)) < 1.0e-4
+    assert np.mean(np.abs(pool5_feat - _pool5_feat)) < 1.0e-4
+    torch._dynamo.reset()
+    logger.info("torch.compile default mode: %s", compile_time)
+    torch._dynamo.reset()
+    rcParam["torch_compile_mode"] = "reduce-overhead"
+    (pool3_feat, pool4_feat, pool5_feat), compile_time = timed(_extract_features)
+    _pool3_feat, _pool4_feat, _pool5_feat = np.load(
+        str(dfbr_features),
+        allow_pickle=True,
+    )
+    assert np.mean(np.abs(pool3_feat - _pool3_feat)) < 1.0e-4
+    assert np.mean(np.abs(pool4_feat - _pool4_feat)) < 1.0e-4
+    assert np.mean(np.abs(pool5_feat - _pool5_feat)) < 1.0e-4
+    torch._dynamo.reset()
+    logger.info("torch.compile reduce-overhead mode: %s", compile_time)
+    rcParam["torch_compile_mode"] = "max-autotune"
+    (pool3_feat, pool4_feat, pool5_feat), compile_time = timed(_extract_features)
+    _pool3_feat, _pool4_feat, _pool5_feat = np.load(
+        str(dfbr_features),
+        allow_pickle=True,
+    )
+    assert np.mean(np.abs(pool3_feat - _pool3_feat)) < 1.0e-4
+    assert np.mean(np.abs(pool4_feat - _pool4_feat)) < 1.0e-4
+    assert np.mean(np.abs(pool5_feat - _pool5_feat)) < 1.0e-4
+    torch._dynamo.reset()
+    rcParam["enable_torch_compile"] = torch_compile_enabled

From a831b7b28bf9a7eb22b70ca2f0d3e4b1d0d66f5f Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 11 Apr 2024 09:51:16 +0000
Subject: [PATCH 04/19] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/test_wsi_registration.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/test_wsi_registration.py b/tests/test_wsi_registration.py
index 470cb281c..cb80fa3fa 100644
--- a/tests/test_wsi_registration.py
+++ b/tests/test_wsi_registration.py
@@ -580,6 +580,7 @@ def test_affine_wsi_transformer(sample_ome_tiff: Path) -> None:
 
         assert np.sum(expected - output) == 0
 
+
 def test_dfbr_feature_extractor_torch_compile(dfbr_features: Path) -> None:
     """Test DFBRFeatureExtractor with torch.compile functionality.
 

From 36239b6e59f59634068f388ebb92699f6ad3e1ef Mon Sep 17 00:00:00 2001
From: Abdol A <u2271662@live.warwick.ac.uk>
Date: Thu, 11 Apr 2024 11:22:55 +0100
Subject: [PATCH 05/19] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20Refactor=20test?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 tests/test_wsi_registration.py | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/tests/test_wsi_registration.py b/tests/test_wsi_registration.py
index cb80fa3fa..9f3b688ef 100644
--- a/tests/test_wsi_registration.py
+++ b/tests/test_wsi_registration.py
@@ -587,10 +587,6 @@ def test_dfbr_feature_extractor_torch_compile(dfbr_features: Path) -> None:
     Args:
         dfbr_features (Path): Path to the expected features.
     """
-    torch_compile_enabled = rcParam["enable_torch_compile"]
-    torch._dynamo.reset()
-    rcParam["enable_torch_compile"] = True
-
     def _extract_features() -> tuple:
         dfbr = DFBRegister()
         fixed_img = np.repeat(
@@ -612,6 +608,9 @@ def _extract_features() -> tuple:
 
         return pool3_feat, pool4_feat, pool5_feat
 
+    torch_compile_enabled = rcParam["enable_torch_compile"]
+    torch._dynamo.reset()
+    rcParam["enable_torch_compile"] = True
     (pool3_feat, pool4_feat, pool5_feat), compile_time = timed(_extract_features)
     _pool3_feat, _pool4_feat, _pool5_feat = np.load(
         str(dfbr_features),
@@ -620,7 +619,6 @@ def _extract_features() -> tuple:
     assert np.mean(np.abs(pool3_feat - _pool3_feat)) < 1.0e-4
     assert np.mean(np.abs(pool4_feat - _pool4_feat)) < 1.0e-4
     assert np.mean(np.abs(pool5_feat - _pool5_feat)) < 1.0e-4
-    torch._dynamo.reset()
     logger.info("torch.compile default mode: %s", compile_time)
     torch._dynamo.reset()
     rcParam["torch_compile_mode"] = "reduce-overhead"
@@ -632,8 +630,8 @@ def _extract_features() -> tuple:
     assert np.mean(np.abs(pool3_feat - _pool3_feat)) < 1.0e-4
     assert np.mean(np.abs(pool4_feat - _pool4_feat)) < 1.0e-4
     assert np.mean(np.abs(pool5_feat - _pool5_feat)) < 1.0e-4
-    torch._dynamo.reset()
     logger.info("torch.compile reduce-overhead mode: %s", compile_time)
+    torch._dynamo.reset()
     rcParam["torch_compile_mode"] = "max-autotune"
     (pool3_feat, pool4_feat, pool5_feat), compile_time = timed(_extract_features)
     _pool3_feat, _pool4_feat, _pool5_feat = np.load(
@@ -643,5 +641,6 @@ def _extract_features() -> tuple:
     assert np.mean(np.abs(pool3_feat - _pool3_feat)) < 1.0e-4
     assert np.mean(np.abs(pool4_feat - _pool4_feat)) < 1.0e-4
     assert np.mean(np.abs(pool5_feat - _pool5_feat)) < 1.0e-4
+    logger.info("torch.compile max-autotune mode: %s", compile_time)
     torch._dynamo.reset()
     rcParam["enable_torch_compile"] = torch_compile_enabled

From 680461c22762e27755688455b1870bd452fd6d89 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 11 Apr 2024 10:23:21 +0000
Subject: [PATCH 06/19] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/test_wsi_registration.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/test_wsi_registration.py b/tests/test_wsi_registration.py
index 9f3b688ef..43cf358c5 100644
--- a/tests/test_wsi_registration.py
+++ b/tests/test_wsi_registration.py
@@ -587,6 +587,7 @@ def test_dfbr_feature_extractor_torch_compile(dfbr_features: Path) -> None:
     Args:
         dfbr_features (Path): Path to the expected features.
     """
+
     def _extract_features() -> tuple:
         dfbr = DFBRegister()
         fixed_img = np.repeat(

From a93edd6cda4fec7e91be40da07dcffa87ce2020a Mon Sep 17 00:00:00 2001
From: Abdol A <u2271662@live.warwick.ac.uk>
Date: Fri, 10 May 2024 15:19:31 +0100
Subject: [PATCH 07/19] =?UTF-8?q?=F0=9F=9A=B8=20Add=20warning=20for=20inco?=
 =?UTF-8?q?mpatible=20GPUs?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 tiatoolbox/models/architecture/utils.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/tiatoolbox/models/architecture/utils.py b/tiatoolbox/models/architecture/utils.py
index 94f970df8..be65163f2 100644
--- a/tiatoolbox/models/architecture/utils.py
+++ b/tiatoolbox/models/architecture/utils.py
@@ -38,6 +38,16 @@ def compile_model(
     if disable:
         return model
 
+    # Check if GPU is compatible with torch.compile
+    if torch.cuda.is_available():
+        device_cap = torch.cuda.get_device_capability()
+        if device_cap not in ((7, 0), (8, 0), (9, 0)):
+            logger.warning(
+                "GPU is not compatible with torch.compile. "
+                "Compatible GPUs include NVIDIA V100, A100, and H100. "
+                "Speedup numbers may be lower than expected."
+            )
+
     # This check will be removed when torch.compile is supported in Python 3.12+
     if sys.version_info >= (3, 12):  # pragma: no cover
         logger.warning(

From cdf530d87a4413118a344c09a1b9ae943eff3170 Mon Sep 17 00:00:00 2001
From: Abdol A <u2271662@live.warwick.ac.uk>
Date: Fri, 10 May 2024 15:48:26 +0100
Subject: [PATCH 08/19] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20=F0=9F=9A=B8=20Remov?=
 =?UTF-8?q?e=20disable=20`torch.compile`=20and=20merge=20with=20options=20?=
 =?UTF-8?q?and=20enable=20by=20default=20if=20compatible?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 tests/models/test_patch_predictor.py          |  5 ++-
 tests/test_wsi_registration.py                |  6 ++--
 tiatoolbox/__init__.py                        | 32 +++++++++++++++----
 tiatoolbox/models/architecture/utils.py       | 20 +++---------
 tiatoolbox/models/engine/patch_predictor.py   |  1 -
 .../tools/registration/wsi_registration.py    |  1 -
 6 files changed, 36 insertions(+), 29 deletions(-)

diff --git a/tests/models/test_patch_predictor.py b/tests/models/test_patch_predictor.py
index 5b94ae189..5fd930138 100644
--- a/tests/models/test_patch_predictor.py
+++ b/tests/models/test_patch_predictor.py
@@ -1247,9 +1247,8 @@ def test_patch_predictor_torch_compile(
         tmp_path (Path): Path to temporary directory.
 
     """
-    torch_compile_enabled = rcParam["enable_torch_compile"]
+    torch_compile_mode = rcParam["torch_compile_mode"]
     torch._dynamo.reset()
-    rcParam["enable_torch_compile"] = True
     rcParam["torch_compile_mode"] = "default"
     _, compile_time = timed(
         test_patch_predictor_api,
@@ -1277,4 +1276,4 @@ def test_patch_predictor_torch_compile(
     )
     logger.info("torch.compile max-autotune mode: %s", compile_time)
     torch._dynamo.reset()
-    rcParam["enable_torch_compile"] = torch_compile_enabled
+    rcParam["torch_compile_mode"] = torch_compile_mode
diff --git a/tests/test_wsi_registration.py b/tests/test_wsi_registration.py
index 43cf358c5..73d3c7653 100644
--- a/tests/test_wsi_registration.py
+++ b/tests/test_wsi_registration.py
@@ -609,9 +609,9 @@ def _extract_features() -> tuple:
 
         return pool3_feat, pool4_feat, pool5_feat
 
-    torch_compile_enabled = rcParam["enable_torch_compile"]
+    torch_compile_mode = rcParam["torch_compile_mode"]
     torch._dynamo.reset()
-    rcParam["enable_torch_compile"] = True
+    rcParam["torch_compile_mode"] = "default"
     (pool3_feat, pool4_feat, pool5_feat), compile_time = timed(_extract_features)
     _pool3_feat, _pool4_feat, _pool5_feat = np.load(
         str(dfbr_features),
@@ -644,4 +644,4 @@ def _extract_features() -> tuple:
     assert np.mean(np.abs(pool5_feat - _pool5_feat)) < 1.0e-4
     logger.info("torch.compile max-autotune mode: %s", compile_time)
     torch._dynamo.reset()
-    rcParam["enable_torch_compile"] = torch_compile_enabled
+    rcParam["torch_compile_mode"] = torch_compile_mode
diff --git a/tiatoolbox/__init__.py b/tiatoolbox/__init__.py
index d82f785b5..93920e83c 100644
--- a/tiatoolbox/__init__.py
+++ b/tiatoolbox/__init__.py
@@ -8,6 +8,7 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, TypedDict
 
+import torch
 import yaml
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -73,10 +74,31 @@ class _RcParam(TypedDict):
 
     TIATOOLBOX_HOME: Path
     pretrained_model_info: dict[str, dict]
-    enable_torch_compile: bool
     torch_compile_mode: str
 
 
+def is_torch_compile_compatible() -> bool:
+    """Check if the current GPU is compatible with torch-compile.
+
+    Returns:
+        bool:
+            True if the GPU is compatible with torch-compile, False
+            otherwise.
+
+    """
+    if torch.cuda.is_available():
+        device_cap = torch.cuda.get_device_capability()
+        if device_cap not in ((7, 0), (8, 0), (9, 0)):
+            logger.warning(
+                "GPU is not compatible with torch.compile. "
+                "Compatible GPUs include NVIDIA V100, A100, and H100. "
+                "Speedup numbers may be lower than expected."
+            )
+            return False
+
+    return True
+
+
 def read_registry_files(path_to_registry: str | Path) -> dict:
     """Reads registry files using importlib_resources.
 
@@ -104,11 +126,9 @@ def read_registry_files(path_to_registry: str | Path) -> dict:
     "pretrained_model_info": read_registry_files(
         "data/pretrained_model.yaml",
     ),  # Load a dictionary of sample files data (names and urls)
-    "enable_torch_compile": False,
-    # Disable `torch-compile`` by default
-    "torch_compile_mode": "default",
-    # Set ``torch-compile`` mode to ``default`` by default
-    # Options: “default”, “reduce-overhead”, “max-autotune”
+    "torch_compile_mode": "default" if is_torch_compile_compatible() else "disable",
+    # Set `torch-compile` mode to `default`if GPU is compatible, otherwise disable
+    # Options: `disable`, `default`, `reduce-overhead`, `max-autotune`
     # or “max-autotune-no-cudagraphs”
 }
 
diff --git a/tiatoolbox/models/architecture/utils.py b/tiatoolbox/models/architecture/utils.py
index be65163f2..5f44f62bc 100644
--- a/tiatoolbox/models/architecture/utils.py
+++ b/tiatoolbox/models/architecture/utils.py
@@ -9,14 +9,13 @@
 import torch
 from torch import nn
 
-from tiatoolbox import logger
+from tiatoolbox import is_torch_compile_compatible, logger
 
 
 def compile_model(
     model: nn.Module | None = None,
     *,
     mode: str = "default",
-    disable: bool = False,
 ) -> Callable:
     """A decorator to compile a model using torch-compile.
 
@@ -25,28 +24,19 @@ def compile_model(
             Model to be compiled.
         mode (str):
             Mode to be used for torch-compile. Available modes are
-            `default`, `reduce-overhead`, `max-autotune`, and
+            `disable`, `default`, `reduce-overhead`, `max-autotune`, and
             `max-autotune-no-cudagraphs`.
-        disable (bool):
-            If True, torch-compile will be disabled.
 
     Returns:
         Callable:
             Compiled model.
 
     """
-    if disable:
+    if mode == "disable":
         return model
 
     # Check if GPU is compatible with torch.compile
-    if torch.cuda.is_available():
-        device_cap = torch.cuda.get_device_capability()
-        if device_cap not in ((7, 0), (8, 0), (9, 0)):
-            logger.warning(
-                "GPU is not compatible with torch.compile. "
-                "Compatible GPUs include NVIDIA V100, A100, and H100. "
-                "Speedup numbers may be lower than expected."
-            )
+    is_torch_compile_compatible()
 
     # This check will be removed when torch.compile is supported in Python 3.12+
     if sys.version_info >= (3, 12):  # pragma: no cover
@@ -64,7 +54,7 @@ def compile_model(
         )
         return model
 
-    return torch.compile(model, mode=mode, disable=disable)
+    return torch.compile(model, mode=mode)
 
 
 def centre_crop(
diff --git a/tiatoolbox/models/engine/patch_predictor.py b/tiatoolbox/models/engine/patch_predictor.py
index 2aede1393..da4420cb0 100644
--- a/tiatoolbox/models/engine/patch_predictor.py
+++ b/tiatoolbox/models/engine/patch_predictor.py
@@ -255,7 +255,6 @@ def __init__(
             compile_model(  # for runtime, such as after wrapping with nn.DataParallel
                 model,
                 mode=rcParam["torch_compile_mode"],
-                disable=not rcParam["enable_torch_compile"],
             )
         )
         self.pretrained_model = pretrained_model
diff --git a/tiatoolbox/tools/registration/wsi_registration.py b/tiatoolbox/tools/registration/wsi_registration.py
index 0906423f8..74e35ffa5 100644
--- a/tiatoolbox/tools/registration/wsi_registration.py
+++ b/tiatoolbox/tools/registration/wsi_registration.py
@@ -342,7 +342,6 @@ def __init__(self: torch.nn.Module) -> None:
         self.pretrained: torch.nn.Sequential = compile_model(
             torchvision.models.vgg16(weights=VGG16_Weights.IMAGENET1K_V1),
             mode=rcParam["torch_compile_mode"],
-            disable=not rcParam["enable_torch_compile"],
         ).features
         self.f_hooks = [
             getattr(self.pretrained, layer).register_forward_hook(

From 663c05b5128d4435eebd8d83c21d2e5527439f52 Mon Sep 17 00:00:00 2001
From: Abdol A <u2271662@live.warwick.ac.uk>
Date: Fri, 10 May 2024 17:25:29 +0100
Subject: [PATCH 09/19] =?UTF-8?q?=F0=9F=9A=B8=20=E2=9C=85=20Add=20test=20f?=
 =?UTF-8?q?or=20`torch.compile`=20and=20extend=20check=20for=20no=20availa?=
 =?UTF-8?q?ble=20cuda?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 tests/test_init.py     | 6 ++++++
 tiatoolbox/__init__.py | 7 +++++++
 2 files changed, 13 insertions(+)

diff --git a/tests/test_init.py b/tests/test_init.py
index 4eddcdd19..01265d84a 100644
--- a/tests/test_init.py
+++ b/tests/test_init.py
@@ -150,3 +150,9 @@ def test_lazy_import_module_not_found() -> None:
             "nonexistent_module",
             Path(__file__).parent.parent / "tiatoolbox",
         )
+
+
+def test_torch_compile_compatibility() -> None:
+    """Test if torch-compile compatibility is checked correctly."""
+    from tiatoolbox import is_torch_compile_compatible
+    assert isinstance(is_torch_compile_compatible(), bool)
\ No newline at end of file
diff --git a/tiatoolbox/__init__.py b/tiatoolbox/__init__.py
index 93920e83c..de36aa02e 100644
--- a/tiatoolbox/__init__.py
+++ b/tiatoolbox/__init__.py
@@ -95,6 +95,13 @@ def is_torch_compile_compatible() -> bool:
                 "Speedup numbers may be lower than expected."
             )
             return False
+    else:
+        logger.warning(
+            "No GPU detected or cuda not installed. "
+            "torch.compile is only supported on selected NVIDIA GPUs. "
+            "Speedup numbers may be lower than expected."
+        )
+        return False
 
     return True
 

From 4f79f1af5c7149a2790c2084480a81649b580673 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 10 May 2024 16:25:53 +0000
Subject: [PATCH 10/19] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/test_init.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/test_init.py b/tests/test_init.py
index 01265d84a..85ddc08d6 100644
--- a/tests/test_init.py
+++ b/tests/test_init.py
@@ -155,4 +155,5 @@ def test_lazy_import_module_not_found() -> None:
 def test_torch_compile_compatibility() -> None:
     """Test if torch-compile compatibility is checked correctly."""
     from tiatoolbox import is_torch_compile_compatible
-    assert isinstance(is_torch_compile_compatible(), bool)
\ No newline at end of file
+
+    assert isinstance(is_torch_compile_compatible(), bool)

From e7304f12191a96ae27f35e9da8174d6b14370f0e Mon Sep 17 00:00:00 2001
From: Abdol A <u2271662@live.warwick.ac.uk>
Date: Fri, 10 May 2024 17:32:02 +0100
Subject: [PATCH 11/19] =?UTF-8?q?=F0=9F=9A=B8=20Extend=20`torch.compile`?=
 =?UTF-8?q?=20for=20next=20gen=20GPUs?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 tests/test_init.py     | 2 +-
 tiatoolbox/__init__.py | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/tests/test_init.py b/tests/test_init.py
index 01265d84a..26529005c 100644
--- a/tests/test_init.py
+++ b/tests/test_init.py
@@ -155,4 +155,4 @@ def test_lazy_import_module_not_found() -> None:
 def test_torch_compile_compatibility() -> None:
     """Test if torch-compile compatibility is checked correctly."""
     from tiatoolbox import is_torch_compile_compatible
-    assert isinstance(is_torch_compile_compatible(), bool)
\ No newline at end of file
+    assert isinstance(is_torch_compile_compatible(), bool)
diff --git a/tiatoolbox/__init__.py b/tiatoolbox/__init__.py
index de36aa02e..b46027d57 100644
--- a/tiatoolbox/__init__.py
+++ b/tiatoolbox/__init__.py
@@ -88,7 +88,8 @@ def is_torch_compile_compatible() -> bool:
     """
     if torch.cuda.is_available():
         device_cap = torch.cuda.get_device_capability()
-        if device_cap not in ((7, 0), (8, 0), (9, 0)):
+        cap_threshold = (7, 0)
+        if device_cap[0] >= cap_threshold[0] and device_cap[1] == cap_threshold[1]:
             logger.warning(
                 "GPU is not compatible with torch.compile. "
                 "Compatible GPUs include NVIDIA V100, A100, and H100. "

From f34f81ceafb9ea18411123a08831dc6f4dda3ff8 Mon Sep 17 00:00:00 2001
From: Abdol A <u2271662@live.warwick.ac.uk>
Date: Mon, 13 May 2024 13:16:20 +0100
Subject: [PATCH 12/19] =?UTF-8?q?=E2=9C=8F=EF=B8=8F=20Fix=20typo=20in=20`t?=
 =?UTF-8?q?orch.compile`=20compatbility=20check?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 tiatoolbox/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tiatoolbox/__init__.py b/tiatoolbox/__init__.py
index b46027d57..2c865ba14 100644
--- a/tiatoolbox/__init__.py
+++ b/tiatoolbox/__init__.py
@@ -98,7 +98,7 @@ def is_torch_compile_compatible() -> bool:
             return False
     else:
         logger.warning(
-            "No GPU detected or cuda not installed. "
+            "No GPU detected or cuda not installed, "
             "torch.compile is only supported on selected NVIDIA GPUs. "
             "Speedup numbers may be lower than expected."
         )

From 52e7e69089f6f103d673639684c7f372b35ab031 Mon Sep 17 00:00:00 2001
From: Abdol A <u2271662@live.warwick.ac.uk>
Date: Mon, 13 May 2024 13:47:57 +0100
Subject: [PATCH 13/19] =?UTF-8?q?=F0=9F=9A=A8=20Potential=20fix=20to=20lin?=
 =?UTF-8?q?ting=20error=20when=20using=20`logger`?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 tiatoolbox/__init__.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tiatoolbox/__init__.py b/tiatoolbox/__init__.py
index 2c865ba14..b93799eb3 100644
--- a/tiatoolbox/__init__.py
+++ b/tiatoolbox/__init__.py
@@ -93,14 +93,16 @@ def is_torch_compile_compatible() -> bool:
             logger.warning(
                 "GPU is not compatible with torch.compile. "
                 "Compatible GPUs include NVIDIA V100, A100, and H100. "
-                "Speedup numbers may be lower than expected."
+                "Speedup numbers may be lower than expected.",
+                stacklevel=2,
             )
             return False
     else:
         logger.warning(
             "No GPU detected or cuda not installed, "
             "torch.compile is only supported on selected NVIDIA GPUs. "
-            "Speedup numbers may be lower than expected."
+            "Speedup numbers may be lower than expected.",
+            stacklevel=2,
         )
         return False
 

From f0306876cb53bb436ddc6580ecfaecc17627f5ae Mon Sep 17 00:00:00 2001
From: Abdol A <u2271662@live.warwick.ac.uk>
Date: Mon, 13 May 2024 15:29:17 +0100
Subject: [PATCH 14/19] =?UTF-8?q?=F0=9F=9A=A8=20Potential=20fix=20to=20lin?=
 =?UTF-8?q?t=20error=20(try=202)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 tiatoolbox/__init__.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tiatoolbox/__init__.py b/tiatoolbox/__init__.py
index b93799eb3..395821b68 100644
--- a/tiatoolbox/__init__.py
+++ b/tiatoolbox/__init__.py
@@ -5,6 +5,7 @@
 import importlib.resources as importlib_resources
 import importlib.util
 import sys
+import warnings
 from pathlib import Path
 from typing import TYPE_CHECKING, TypedDict
 
@@ -90,7 +91,7 @@ def is_torch_compile_compatible() -> bool:
         device_cap = torch.cuda.get_device_capability()
         cap_threshold = (7, 0)
         if device_cap[0] >= cap_threshold[0] and device_cap[1] == cap_threshold[1]:
-            logger.warning(
+            warnings.warn(
                 "GPU is not compatible with torch.compile. "
                 "Compatible GPUs include NVIDIA V100, A100, and H100. "
                 "Speedup numbers may be lower than expected.",
@@ -98,7 +99,7 @@ def is_torch_compile_compatible() -> bool:
             )
             return False
     else:
-        logger.warning(
+        warnings.warn(
             "No GPU detected or cuda not installed, "
             "torch.compile is only supported on selected NVIDIA GPUs. "
             "Speedup numbers may be lower than expected.",

From 43881b121c95344907f10e8de9990af159bd2dda Mon Sep 17 00:00:00 2001
From: Abdol A <u2271662@live.warwick.ac.uk>
Date: Mon, 13 May 2024 16:51:28 +0100
Subject: [PATCH 15/19] =?UTF-8?q?=F0=9F=90=9B=20Move=20`torch.compile`=20c?=
 =?UTF-8?q?omptability=20check=20function=20to=20avoid=20`test=5Flogger=5F?=
 =?UTF-8?q?output`=20error?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 tests/test_init.py                      |  2 +-
 tiatoolbox/__init__.py                  | 38 ++-----------------------
 tiatoolbox/models/architecture/utils.py | 34 +++++++++++++++++++++-
 3 files changed, 36 insertions(+), 38 deletions(-)

diff --git a/tests/test_init.py b/tests/test_init.py
index 85ddc08d6..9c8b0c5cf 100644
--- a/tests/test_init.py
+++ b/tests/test_init.py
@@ -154,6 +154,6 @@ def test_lazy_import_module_not_found() -> None:
 
 def test_torch_compile_compatibility() -> None:
     """Test if torch-compile compatibility is checked correctly."""
-    from tiatoolbox import is_torch_compile_compatible
+    from tiatoolbox.models.architecture.utils import is_torch_compile_compatible
 
     assert isinstance(is_torch_compile_compatible(), bool)
diff --git a/tiatoolbox/__init__.py b/tiatoolbox/__init__.py
index 395821b68..80ad1b4ff 100644
--- a/tiatoolbox/__init__.py
+++ b/tiatoolbox/__init__.py
@@ -5,11 +5,9 @@
 import importlib.resources as importlib_resources
 import importlib.util
 import sys
-import warnings
 from pathlib import Path
 from typing import TYPE_CHECKING, TypedDict
 
-import torch
 import yaml
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -78,38 +76,6 @@ class _RcParam(TypedDict):
     torch_compile_mode: str
 
 
-def is_torch_compile_compatible() -> bool:
-    """Check if the current GPU is compatible with torch-compile.
-
-    Returns:
-        bool:
-            True if the GPU is compatible with torch-compile, False
-            otherwise.
-
-    """
-    if torch.cuda.is_available():
-        device_cap = torch.cuda.get_device_capability()
-        cap_threshold = (7, 0)
-        if device_cap[0] >= cap_threshold[0] and device_cap[1] == cap_threshold[1]:
-            warnings.warn(
-                "GPU is not compatible with torch.compile. "
-                "Compatible GPUs include NVIDIA V100, A100, and H100. "
-                "Speedup numbers may be lower than expected.",
-                stacklevel=2,
-            )
-            return False
-    else:
-        warnings.warn(
-            "No GPU detected or cuda not installed, "
-            "torch.compile is only supported on selected NVIDIA GPUs. "
-            "Speedup numbers may be lower than expected.",
-            stacklevel=2,
-        )
-        return False
-
-    return True
-
-
 def read_registry_files(path_to_registry: str | Path) -> dict:
     """Reads registry files using importlib_resources.
 
@@ -137,8 +103,8 @@ def read_registry_files(path_to_registry: str | Path) -> dict:
     "pretrained_model_info": read_registry_files(
         "data/pretrained_model.yaml",
     ),  # Load a dictionary of sample files data (names and urls)
-    "torch_compile_mode": "default" if is_torch_compile_compatible() else "disable",
-    # Set `torch-compile` mode to `default`if GPU is compatible, otherwise disable
+    "torch_compile_mode": "default",
+    # Set `torch-compile` mode to `default`
     # Options: `disable`, `default`, `reduce-overhead`, `max-autotune`
     # or “max-autotune-no-cudagraphs”
 }
diff --git a/tiatoolbox/models/architecture/utils.py b/tiatoolbox/models/architecture/utils.py
index 5f44f62bc..5f6b5da18 100644
--- a/tiatoolbox/models/architecture/utils.py
+++ b/tiatoolbox/models/architecture/utils.py
@@ -9,7 +9,39 @@
 import torch
 from torch import nn
 
-from tiatoolbox import is_torch_compile_compatible, logger
+from tiatoolbox import logger
+
+
+def is_torch_compile_compatible() -> bool:
+    """Check if the current GPU is compatible with torch-compile.
+
+    Returns:
+        bool:
+            True if the GPU is compatible with torch-compile, False
+            otherwise.
+
+    """
+    if torch.cuda.is_available():
+        device_cap = torch.cuda.get_device_capability()
+        cap_threshold = (7, 0)
+        if device_cap[0] >= cap_threshold[0] and device_cap[1] == cap_threshold[1]:
+            logger.warning(
+                "GPU is not compatible with torch.compile. "
+                "Compatible GPUs include NVIDIA V100, A100, and H100. "
+                "Speedup numbers may be lower than expected.",
+                stacklevel=2,
+            )
+            return False
+    else:
+        logger.warning(
+            "No GPU detected or cuda not installed, "
+            "torch.compile is only supported on selected NVIDIA GPUs. "
+            "Speedup numbers may be lower than expected.",
+            stacklevel=2,
+        )
+        return False
+
+    return True
 
 
 def compile_model(

From 3d315682c3bf1f03cfcce8b206d1b8ee2966fcd3 Mon Sep 17 00:00:00 2001
From: Abdol A <u2271662@live.warwick.ac.uk>
Date: Mon, 13 May 2024 17:46:53 +0100
Subject: [PATCH 16/19] =?UTF-8?q?=E2=9C=85=20Move=20test=20to=20`test=5Fut?=
 =?UTF-8?q?ils`?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 tests/test_init.py             | 7 -------
 tests/test_utils.py            | 7 +++++++
 tests/test_wsi_registration.py | 1 +
 3 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/tests/test_init.py b/tests/test_init.py
index 9c8b0c5cf..4eddcdd19 100644
--- a/tests/test_init.py
+++ b/tests/test_init.py
@@ -150,10 +150,3 @@ def test_lazy_import_module_not_found() -> None:
             "nonexistent_module",
             Path(__file__).parent.parent / "tiatoolbox",
         )
-
-
-def test_torch_compile_compatibility() -> None:
-    """Test if torch-compile compatibility is checked correctly."""
-    from tiatoolbox.models.architecture.utils import is_torch_compile_compatible
-
-    assert isinstance(is_torch_compile_compatible(), bool)
diff --git a/tests/test_utils.py b/tests/test_utils.py
index cf76028aa..44377098f 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -1837,3 +1837,10 @@ def test_torch_compile_already_compiled() -> None:
     # Check that the recompiled model
     # is the same as the original compiled model
     assert compiled_model == recompiled_model
+
+
+def test_torch_compile_compatibility() -> None:
+    """Test if torch-compile compatibility is checked correctly."""
+    from tiatoolbox.models.architecture.utils import is_torch_compile_compatible
+
+    assert isinstance(is_torch_compile_compatible(), bool)
diff --git a/tests/test_wsi_registration.py b/tests/test_wsi_registration.py
index 73d3c7653..4e0f07366 100644
--- a/tests/test_wsi_registration.py
+++ b/tests/test_wsi_registration.py
@@ -586,6 +586,7 @@ def test_dfbr_feature_extractor_torch_compile(dfbr_features: Path) -> None:
 
     Args:
         dfbr_features (Path): Path to the expected features.
+
     """
 
     def _extract_features() -> tuple:

From e8335a88718aca4d373d4ba9830fe7d99651a946 Mon Sep 17 00:00:00 2001
From: Abdol A <u2271662@live.warwick.ac.uk>
Date: Mon, 13 May 2024 18:52:09 +0100
Subject: [PATCH 17/19] =?UTF-8?q?=F0=9F=90=9B=20=E2=9C=85=20Update=20tests?=
 =?UTF-8?q?=20for=20coverage=20and=20fix=20compatibility=20check=20conditi?=
 =?UTF-8?q?on=20statement?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 tests/test_utils.py                     | 27 ++++++++++++++++---------
 tiatoolbox/models/architecture/utils.py |  7 +++----
 2 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/tests/test_utils.py b/tests/test_utils.py
index 44377098f..b1ca5a2da 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -19,7 +19,7 @@
 from shapely.geometry import Polygon
 
 from tests.test_annotation_stores import cell_polygon
-from tiatoolbox import utils
+from tiatoolbox import rcParam, utils
 from tiatoolbox.annotation.storage import SQLiteStore
 from tiatoolbox.models.architecture import fetch_pretrained_weights
 from tiatoolbox.models.architecture.utils import compile_model
@@ -1825,18 +1825,25 @@ def test_patch_pred_store_persist_ext(tmp_path: pytest.TempPathFactory) -> None:
 
 def test_torch_compile_already_compiled() -> None:
     """Test that torch_compile does not recompile a model that is already compiled."""
-    # Create a simple model
+    torch_compile_modes = [
+        "disable",
+        "default",
+        "reduce-overhead",
+        "max-autotune",
+        "max-autotune-no-cudagraphs",
+    ]
+    current_torch_compile_mode = rcParam["torch_compile_mode"]
     model = torch.nn.Sequential(torch.nn.Linear(10, 10), torch.nn.Linear(10, 10))
 
-    # Compile the model
-    compiled_model = compile_model(model)
-
-    # Compile the model again
-    recompiled_model = compile_model(compiled_model)
+    for mode in torch_compile_modes:
+        torch._dynamo.reset()
+        rcParam["torch_compile_mode"] = mode
+        compiled_model = compile_model(model)
+        recompiled_model = compile_model(compiled_model)
+        assert compiled_model == recompiled_model
 
-    # Check that the recompiled model
-    # is the same as the original compiled model
-    assert compiled_model == recompiled_model
+    torch._dynamo.reset()
+    rcParam["torch_compile_mode"] = current_torch_compile_mode
 
 
 def test_torch_compile_compatibility() -> None:
diff --git a/tiatoolbox/models/architecture/utils.py b/tiatoolbox/models/architecture/utils.py
index 5f6b5da18..2150c31a1 100644
--- a/tiatoolbox/models/architecture/utils.py
+++ b/tiatoolbox/models/architecture/utils.py
@@ -21,10 +21,9 @@ def is_torch_compile_compatible() -> bool:
             otherwise.
 
     """
-    if torch.cuda.is_available():
+    if torch.cuda.is_available():  # pragma: no cover
         device_cap = torch.cuda.get_device_capability()
-        cap_threshold = (7, 0)
-        if device_cap[0] >= cap_threshold[0] and device_cap[1] == cap_threshold[1]:
+        if device_cap not in ((7, 0), (8, 0), (9, 0)):
             logger.warning(
                 "GPU is not compatible with torch.compile. "
                 "Compatible GPUs include NVIDIA V100, A100, and H100. "
@@ -41,7 +40,7 @@ def is_torch_compile_compatible() -> bool:
         )
         return False
 
-    return True
+    return True  # pragma: no cover
 
 
 def compile_model(

From 8aa34588346972d1ddb790dd90325ff8ab962852 Mon Sep 17 00:00:00 2001
From: Abdol A <u2271662@live.warwick.ac.uk>
Date: Tue, 14 May 2024 10:26:51 +0100
Subject: [PATCH 18/19] =?UTF-8?q?=F0=9F=90=9B=20Fix=20a=20bug=20where=20`t?=
 =?UTF-8?q?orch.compile`=20mode=20does=20not=20change=20in=20a=20test?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 tests/test_utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/test_utils.py b/tests/test_utils.py
index b1ca5a2da..16c62e382 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -1838,8 +1838,8 @@ def test_torch_compile_already_compiled() -> None:
     for mode in torch_compile_modes:
         torch._dynamo.reset()
         rcParam["torch_compile_mode"] = mode
-        compiled_model = compile_model(model)
-        recompiled_model = compile_model(compiled_model)
+        compiled_model = compile_model(model, mode=mode)
+        recompiled_model = compile_model(compiled_model, mode=mode)
         assert compiled_model == recompiled_model
 
     torch._dynamo.reset()

From f054d15d77d3c5f8a481dbc88ccff91bc00a3d84 Mon Sep 17 00:00:00 2001
From: Abdol A <u2271662@live.warwick.ac.uk>
Date: Tue, 14 May 2024 10:57:38 +0100
Subject: [PATCH 19/19] =?UTF-8?q?=E2=9C=85=20Add=20a=20seperate=20test=20f?=
 =?UTF-8?q?or=20disabling=20`torch.compile`?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 tests/test_utils.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/tests/test_utils.py b/tests/test_utils.py
index 16c62e382..d35ae31e5 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -1826,7 +1826,6 @@ def test_patch_pred_store_persist_ext(tmp_path: pytest.TempPathFactory) -> None:
 def test_torch_compile_already_compiled() -> None:
     """Test that torch_compile does not recompile a model that is already compiled."""
     torch_compile_modes = [
-        "disable",
         "default",
         "reduce-overhead",
         "max-autotune",
@@ -1846,6 +1845,13 @@ def test_torch_compile_already_compiled() -> None:
     rcParam["torch_compile_mode"] = current_torch_compile_mode
 
 
+def test_torch_compile_disable() -> None:
+    """Test torch_compile's disable mode."""
+    model = torch.nn.Sequential(torch.nn.Linear(10, 10), torch.nn.Linear(10, 10))
+    compiled_model = compile_model(model, mode="disable")
+    assert model == compiled_model
+
+
 def test_torch_compile_compatibility() -> None:
     """Test if torch-compile compatibility is checked correctly."""
     from tiatoolbox.models.architecture.utils import is_torch_compile_compatible