Merge pull request #1433 from bghira/bugfix/vae-cache-clear-meta-device

vae cache clear at epoch flip has meta tensor error

Author: bghira

helpers/caching/vae.py

@@ -288,16 +288,7 @@ def init_vae(self):
             if args.pretrained_vae_model_name_or_path is None
             else args.pretrained_vae_model_name_or_path
         )
-        precached_vae = StateTracker.get_vae()
-        self.vae = precached_vae or AutoencoderClass.from_pretrained(
-            vae_path,
-            subfolder="vae" if args.pretrained_vae_model_name_or_path is None else None,
-            revision=args.revision,
-            force_upcast=False,
-        ).to(self.accelerator.device)
-        if self.vae.device != self.accelerator.device:
-            self.vae = self.vae.to(self.accelerator.device)
-        StateTracker.set_vae(self.vae)
+        self.vae = self.model.get_vae()
 
     def rebuild_cache(self):
         """