Update autoencoder_kl as well

microsoft · Jan 8, 2024 · 683a732 · 683a732
1 parent 68dcb3c
commit 683a732
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepspeed/module_inject/containers/vae.py b/deepspeed/module_inject/containers/vae.py
@@ -17,7 +17,7 @@ def __init__(self):
                 self._orig_layer_class = diffusers.models.autoencoders.vae.AutoencoderKL
             else:
                 # Diffusers >= 0.12.0 changes location of AutoencoderKL
-                self._orig_layer_class = diffusers.models.autoencoder_kl.AutoencoderKL
+                self._orig_layer_class = diffusers.models.autoencoders.autoencoder_kl.AutoencoderKL
         except ImportError:
             self._orig_layer_class = None