INNER CODE UNIT · Python

forward

lucidrains/vector-quantize-pytorch · vector_quantize_pytorch/evo_lookup_free_quantization.py:74

    def forward(self, x, **kwargs):
        latents = self.encoder(x)
        is_2d = latents.ndim == 2

        if is_2d:
            latents = rearrange(latents, 'b d -> b 1 d')

        quantized, indices, aux_loss = self.lfq(latents, **kwargs)

        if is_2d:
            quantized = rearrange(quantized, 'b 1 d -> b d')
            indices = rearrange(indices, 'b 1 ... -> b ...')

        reconstructed = self.decoder(quantized)
        return Return(reconstructed, indices, aux_loss)

    @torch.no_grad()
    def encode(self, x, return_signs = False):

View source record →