1 files changed, 3 insertions, 0 deletions

diff --git a/training/functional.py b/training/functional.py
index 68ea40c..38dd59f 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -700,8 +700,11 @@ def train(
     vae.requires_grad_(False)
     vae.eval()
 
+    vae = torch.compile(vae, backend='hidet')
+
     if compile_unet:
         unet = torch.compile(unet, backend='hidet')
+        # unet = torch.compile(unet)
 
     callbacks = strategy.callbacks(
         accelerator=accelerator,