Method forward

models/light_decoder.py:112–134

Args: z: Latent tensor of shape [B, in_channels, H', W'] Returns: Decoded image of shape [B, out_channels, H, W]

(self, z: torch.Tensor)

        )

    def forward(self, z: torch.Tensor) -> torch.Tensor:
        """
        Args:
            z: Latent tensor of shape [B, in_channels, H', W']
        Returns:
            Decoded image of shape [B, out_channels, H, W]
        """
        # conv_in
        h = self.conv_in(z)

        # mid block
        h = self.mid_block(h)

        # up blocks
        for up_block in self.up_blocks:
            h = up_block(h)

        # output
        h = self.conv_norm_out(h)
        h = self.conv_act(h)
        h = self.conv_out(h)

        return h