1 files changed, 4 insertions, 2 deletions

diff --git a/train_ti.py b/train_ti.py
index d1defb3..7d10317 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -538,8 +538,10 @@ def main():
     tokenizer.set_dropout(args.vector_dropout)
 
     vae.enable_slicing()
-    vae.set_use_memory_efficient_attention_xformers(True)
-    unet.enable_xformers_memory_efficient_attention()
+    # vae.set_use_memory_efficient_attention_xformers(True)
+    # unet.enable_xformers_memory_efficient_attention()
+
+    # unet = torch.compile(unet)
 
     if args.gradient_checkpointing:
         unet.enable_gradient_checkpointing()