1 files changed, 2 insertions, 2 deletions

diff --git a/training/functional.py b/training/functional.py
index 10560e5..fd3f9f4 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -710,8 +710,8 @@ def train(
     vae = torch.compile(vae, backend='hidet')
 
     if compile_unet:
-        # unet = torch.compile(unet, backend='hidet')
-        unet = torch.compile(unet, mode="reduce-overhead")
+        unet = torch.compile(unet, backend='hidet')
+        # unet = torch.compile(unet, mode="reduce-overhead")
 
     callbacks = strategy.callbacks(
         accelerator=accelerator,