def METHOD_NAME(self, x, produce_targets=False):
    result = {"num_vars": self.num_vars}
    if self.time_first:
        x = x.transpose(1, 2)
    bsz, fsz, tsz = x.shape
    ze = self.projection(x)
    ze_ = ze.view(bsz, self.groups, self.var_dim, tsz).permute(0, 3, 1, 2)
    d = (
        (ze_.unsqueeze(0) - self.expand_embedding.unsqueeze(1).unsqueeze(1))
        .view(self.num_vars, bsz, tsz, self.groups, -1)
        .norm(dim=-1, p=2)
    )
    idx = d.argmin(dim=0)
    zq = (
        torch.stack(
            [
                self.expand_embedding[idx[..., group], group]
                for group in range(self.groups)
            ],
            dim=-2,
        )
        .view(bsz, tsz, self.groups * self.var_dim)
        .permute(0, 2, 1)
    )
    assert ze.shape == zq.shape, (ze.shape, zq.shape)
    x = self._pass_grad(ze, zq)
    hard_x = (
        idx.new_zeros(bsz * tsz * self.groups, self.num_vars)
        .scatter_(-1, idx.view(-1, 1), 1.0)
        .view(bsz * tsz, self.groups, -1)
    )
    hard_probs = torch.mean(hard_x.float(), dim=0)
    result["code_perplexity"] = torch.exp(
        -torch.sum(hard_probs * torch.log(hard_probs + 1e-7), dim=-1)
    ).sum()
    if produce_targets:
        result["targets"] = idx
    if self.time_first:
        x = x.transpose(1, 2)  # BCT -> BTC
    result["x"] = x
    ze = ze.float()
    zq = zq.float()
    latent_loss = self.mse_mean(zq, ze.detach())
    commitment_loss = self.mse_mean(ze, zq.detach())
    result["kmeans_loss"] = latent_loss + self.gamma * commitment_loss
    return result