From 74a5974ba30c170198890e59c92463bf5319fe64 Mon Sep 17 00:00:00 2001
From: Volpeon <git@volpeon.ink>
Date: Sat, 29 Apr 2023 16:35:41 +0200
Subject: torch.compile won't work yet, keep code prepared

---
 training/functional.py | 3 +++
 1 file changed, 3 insertions(+)

(limited to 'training/functional.py')

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,
-- 
cgit v1.2.3-70-g09d2