fix MIC-DKFZ#2219 by disabling torch.compile if the device is CPU

mzenk · May 28, 2024 · 5549ead · 5549ead
1 parent 7141d4f
commit 5549ead
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/nnunetv2/training/nnUNetTrainer/nnUNetTrainer.py b/nnunetv2/training/nnUNetTrainer/nnUNetTrainer.py
@@ -237,6 +237,8 @@ def initialize(self):
 
     def _do_i_compile(self):
         # new default: compile is enabled!
+        if self.device == torch.device('cpu'):
+            return False
         if 'nnUNet_compile' not in os.environ.keys():
             return True
         else: