INNER CODE UNIT · Python

encode

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

    def encode(self, x, return_signs = False):
        with temp_eval(self):
            latents = self.encoder(x)
            is_2d = latents.ndim == 2

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

            quantized, indices, _ = self.lfq(latents)

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

            if return_signs:
                return torch.where(quantized > 0, 1.0, -1.0)

            return (quantized > 0).float()

View source record →