Avoid model recompilation due to varying prompt lengths

1 files changed, 2 insertions, 1 deletions

diff --git a/train_lora.py b/train_lora.py
index a58bef7..12d7e72 100644
--- a/train_lora.py
+++ b/train_lora.py
@@ -49,7 +49,7 @@ torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.benchmark = True
 
 torch._dynamo.config.log_level = logging.WARNING
-torch._dynamo.config.suppress_errors = True
+# torch._dynamo.config.suppress_errors = True
 
 hidet.torch.dynamo_config.use_tensor_core(True)
 hidet.torch.dynamo_config.search_space(0)
@@ -992,6 +992,7 @@ def main():
         VlpnDataModule,
         data_file=args.train_data_file,
         tokenizer=tokenizer,
+        constant_prompt_length=args.compile_unet,
         class_subdir=args.class_image_dir,
         with_guidance=args.guidance_scale != 0,
         num_class_images=args.num_class_images,