Method encode

pipelines/modeling_lightweight_vae.py:183–192

(
        self, x: torch.Tensor, return_dict: bool = True
    )

    @apply_forward_hook
    def encode(
        self, x: torch.Tensor, return_dict: bool = True
    ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
        moments = self.encoder(x)
        posterior = DiagonalGaussianDistribution(moments)

        if not return_dict:
            return (posterior,)

        return AutoencoderKLOutput(latent_dist=posterior)

    @apply_forward_hook
    def decode(self, z: torch.Tensor, interpolate=True, return_dict: bool = True) -> Union[DecoderOutput, Tuple[torch.Tensor]]: