def METHOD_NAME(self, x, produce_targets=False):
    result = {"num_vars": self.num_vars * self.groups}
    if not self.time_first:
        x = x.transpose(1, 2)
    bsz, tsz, fsz = x.shape
    x = x.reshape(-1, fsz)
    x = self.weight_proj(x)
    x = x.view(bsz * tsz * self.groups, -1)
    _, k = x.max(-1)
    hard_x = (
        x.new_zeros(*x.shape)
        .scatter_(-1, k.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()
    avg_probs = torch.softmax(
        x.view(bsz * tsz, self.groups, -1).float(), dim=-1
    ).mean(dim=0)
    result["prob_perplexity"] = torch.exp(
        -torch.sum(avg_probs * torch.log(avg_probs + 1e-7), dim=-1)
    ).sum()
    result["temp"] = self.curr_temp
    if self.training:
        x = F.gumbel_softmax(x.float(), tau=self.curr_temp, hard=True).type_as(x)
    else:
        x = hard_x
    x = x.view(bsz * tsz, -1)
    vars = self.vars
    if self.combine_groups:
        vars = vars.repeat(1, self.groups, 1)
    if produce_targets:
        result["targets"] = (
            x.view(bsz * tsz * self.groups, -1)
            .argmax(dim=-1)
            .view(bsz, tsz, self.groups)
            .detach()
        )
    x = x.unsqueeze(-1) * vars
    x = x.view(bsz * tsz, self.groups, self.num_vars, -1)
    x = x.sum(-2)
    x = x.view(bsz, tsz, -1)
    if not self.time_first:
        x = x.transpose(1, 2)  # BTC -> BCT
    result["x"] = x
    return result