Unverified Commit 6529ee67 authored by Vinh H. Pham's avatar Vinh H. Pham Committed by GitHub
Browse files

Make VAE compatible to torch.compile() (#7984)



make VAE compatible to torch.compile()
Co-authored-by: default avatarYiYi Xu <yixu310@gmail.com>
parent df2bc5ef
...@@ -300,7 +300,7 @@ class AutoencoderKL(ModelMixin, ConfigMixin, FromOriginalModelMixin): ...@@ -300,7 +300,7 @@ class AutoencoderKL(ModelMixin, ConfigMixin, FromOriginalModelMixin):
decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)] decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]
decoded = torch.cat(decoded_slices) decoded = torch.cat(decoded_slices)
else: else:
decoded = self._decode(z).sample decoded = self._decode(z, return_dict=False)[0]
if not return_dict: if not return_dict:
return (decoded,) return (decoded,)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment