From c7bfb8b22a4caaa5a32c503d7ff3f6db61390af0 Mon Sep 17 00:00:00 2001
From: "Wang, Yi" <yi.a.wang@intel.com>
Date: Mon, 27 Nov 2023 19:43:49 +0800
Subject: [PATCH] set the model to train state before accelerator prepare
 (#5099)

Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>
---
 examples/textual_inversion/textual_inversion.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/examples/textual_inversion/textual_inversion.py b/examples/textual_inversion/textual_inversion.py
index 8e932add92..7fea4fdb64 100644
--- a/examples/textual_inversion/textual_inversion.py
+++ b/examples/textual_inversion/textual_inversion.py
@@ -761,6 +761,7 @@ def main():
         num_cycles=args.lr_num_cycles,
     )
 
+    text_encoder.train()
     # Prepare everything with our `accelerator`.
     text_encoder, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
         text_encoder, optimizer, train_dataloader, lr_scheduler