From 74a5974ba30c170198890e59c92463bf5319fe64 Mon Sep 17 00:00:00 2001 From: Volpeon 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') 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