1 files changed, 3 insertions, 3 deletions

diff --git a/training/functional.py b/training/functional.py
index e7cc20f..68ea40c 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -672,6 +672,7 @@ def train(
     optimizer: torch.optim.Optimizer,
     lr_scheduler: torch.optim.lr_scheduler._LRScheduler,
     strategy: TrainingStrategy,
+    compile_unet: bool = False,
     no_val: bool = False,
     num_train_epochs: int = 100,
     gradient_accumulation_steps: int = 1,
@@ -699,9 +700,8 @@ def train(
     vae.requires_grad_(False)
     vae.eval()
 
-    unet = torch.compile(unet)
-    text_encoder = torch.compile(text_encoder)
-    vae = torch.compile(vae)
+    if compile_unet:
+        unet = torch.compile(unet, backend='hidet')
 
     callbacks = strategy.callbacks(
         accelerator=accelerator,