🚧 Add torch.compile to SemanticSegmentor

TissueImageAnalytics · Sep 26, 2024 · ba1776e · ba1776e
1 parent 8cc2fb4
commit ba1776e
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/tiatoolbox/models/engine/semantic_segmentor.py b/tiatoolbox/models/engine/semantic_segmentor.py
@@ -17,8 +17,9 @@
 import torch.utils.data as torch_data
 import tqdm
 
-from tiatoolbox import logger
+from tiatoolbox import logger, rcParam
 from tiatoolbox.models.architecture import get_pretrained_model
+from tiatoolbox.models.architecture.utils import compile_model
 from tiatoolbox.models.models_abc import IOConfigABC
 from tiatoolbox.tools.patchextraction import PatchExtractor
 from tiatoolbox.utils import imread, misc
@@ -563,7 +564,12 @@ def __init__(
         self.masks = None
 
         self.dataset_class: WSIStreamDataset = dataset_class
-        self.model = model
+        self.model = (
+            compile_model(
+                model,
+                mode=rcParam["torch_compile_mode"],
+            )
+        )
         self.pretrained_model = pretrained_model
         self.batch_size = batch_size
         self.num_loader_workers = num_loader_workers