Fix issue that left text encoder on gpu after caching text embeddings.

jobs/process/BaseSDTrainProcess.py

@@ -2254,6 +2254,16 @@ class BaseSDTrainProcess(BaseTrainProcess):
                 pool_datasets=self.train_config.pool_datasets,
             )
 
+        # a fresh text-embedding cache leaves the text encoder on the gpu (the
+        # datasets don't restore placement: that would bounce it per dataset).
+        # Park it once here, before accelerator.prepare brings the unet up.
+        if (
+            self.is_caching_text_embeddings
+            and self.model_config.low_vram
+            and self.sd.text_encoder is not None
+        ):
+            self.sd.text_encoder_to('cpu')
+
         flush()
         self.last_save_step = self.step_num
         ### HOOK ###