//  [python]
def forward(self,
                z,
                segm_map,
                temp=None,
                rescale_logits=False,
                return_logits=False):
        assert temp is None or temp == 1.0, "Only for interface compatible with Gumbel"
        assert rescale_logits == False, "Only for interface compatible with Gumbel"
        assert return_logits == False, "Only for interface compatible with Gumbel"

        segm_map = F.interpolate(segm_map, size=z.size()[2:], mode='nearest')
        # reshape z -> (batch, height, width, channel) and flatten
        z = rearrange(z, 'b c h w -> b h w c').contiguous()
        z_flattened = z.view(-1, self.e_dim)

        # flatten segm_map (b, h, w)
        segm_map_flatten = segm_map.view(-1)

        z_q = torch.zeros_like(z_flattened)
        min_encoding_indices_list = []
        min_encoding_indices_continual = torch.full(
            segm_map_flatten.size(),
            fill_value=-1,
            dtype=torch.long,
            device=segm_map_flatten.device)
        for codebook_idx in range(18):
            min_encoding_indices = torch.full(
                segm_map_flatten.size(),
                fill_value=-1,
                dtype=torch.long,
                device=segm_map_flatten.device)
            if torch.sum(segm_map_flatten == codebook_idx) > 0:
                z_selected = z_flattened[segm_map_flatten == codebook_idx]
                # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
                d_selected = torch.sum(
                    z_selected**2, dim=1, keepdim=True) + torch.sum(
                        self.embedding_list[codebook_idx].weight**2,
                        dim=1) - 2 * torch.einsum(
                            'bd,dn->bn', z_selected,
                            rearrange(self.embedding_list[codebook_idx].weight,
                                      'n d -> d n'))
                min_encoding_indices_selected = torch.argmin(d_selected, dim=1)
                z_q_selected = self.embedding_list[codebook_idx](
                    min_encoding_indices_selected)
                z_q[segm_map_flatten == codebook_idx] = z_q_selected
                min_encoding_indices[
                    segm_map_flatten ==
                    codebook_idx] = min_encoding_indices_selected
                min_encoding_indices_continual[
                    segm_map_flatten ==
                    codebook_idx] = min_encoding_indices_selected + 1024 * codebook_idx
            min_encoding_indices = min_encoding_indices.reshape(
                z.shape[0], z.shape[1], z.shape[2])
            min_encoding_indices_list.append(min_encoding_indices)

        min_encoding_indices_continual = min_encoding_indices_continual.reshape(
            z.shape[0], z.shape[1], z.shape[2])
        z_q = z_q.view(z.shape)
        perplexity = None

        # compute loss for embedding
        if not self.legacy:
            loss = self.beta * torch.mean((z_q.detach()-z)**2) + \
                   torch.mean((z_q - z.detach()) ** 2)
        else:
            loss = torch.mean((z_q.detach()-z)**2) + self.beta * \
                   torch.mean((z_q - z.detach()) ** 2)

        # preserve gradients
        z_q = z + (z_q - z).detach()

        # reshape back to match original input shape
        z_q = rearrange(z_q, 'b h w c -> b c h w').contiguous()

        return z_q, loss, (perplexity, min_encoding_indices_continual,
                           min_encoding_indices_list)