Method encode

diff2flow/tiny_autoencoder.py:66–73

Args: x: torch.Tensor, shape (b, 3, h, w) in [-1, 1]

(self, x)

    @torch.no_grad()
    def encode(self, x):
        """
        Args:
            x: torch.Tensor, shape (b, 3, h, w) in [-1, 1]
        """
        # scale to [0, 1]
        x = x.div(2).add(0.5)
        return self.encoder(x)
    
    @torch.no_grad()
    def decode(self, z):