1 files changed, 6 insertions, 0 deletions

diff --git a/train_lora.py b/train_lora.py
index 64346bc..74afeed 100644
--- a/train_lora.py
+++ b/train_lora.py
@@ -456,6 +456,11 @@ def parse_args():
         ),
     )
     parser.add_argument(
+        "--compile_unet",
+        action="store_true",
+        help="Compile UNet with Torch Dynamo.",
+    )
+    parser.add_argument(
         "--lora_rank",
         type=int,
         default=256,
@@ -892,6 +897,7 @@ def main():
         noise_scheduler=noise_scheduler,
         dtype=weight_dtype,
         seed=args.seed,
+        compile_unet=args.compile_unet,
         guidance_scale=args.guidance_scale,
         prior_loss_weight=args.prior_loss_weight if args.num_class_images != 0 else 0,
         sample_scheduler=sample_scheduler,