From d12a0c1c02c18b23918fe3d0cf94f56086ef3e85 Mon Sep 17 00:00:00 2001 From: Fabian Isensee Date: Tue, 28 May 2024 09:46:13 +0200 Subject: [PATCH] disable torch.compile for windows #2219 --- nnunetv2/training/nnUNetTrainer/nnUNetTrainer.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/nnunetv2/training/nnUNetTrainer/nnUNetTrainer.py b/nnunetv2/training/nnUNetTrainer/nnUNetTrainer.py index 4c179f172..d3803f391 100644 --- a/nnunetv2/training/nnUNetTrainer/nnUNetTrainer.py +++ b/nnunetv2/training/nnUNetTrainer/nnUNetTrainer.py @@ -237,8 +237,16 @@ def initialize(self): def _do_i_compile(self): # new default: compile is enabled! + + # CPU compile crashes for 2D models. Not sure if we even want to support CPU compile!? Better disable if self.device == torch.device('cpu'): return False + + # default torch.compile doesn't work on windows because there are apparently no triton wheels for it + # https://discuss.pytorch.org/t/windows-support-timeline-for-torch-compile/182268/2 + if os.name == 'nt': + return False + if 'nnUNet_compile' not in os.environ.keys(): return True else: