def METHOD_NAME(self, x, return_ids=False):
    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)
    # Calculate quantize prob perplexity
    num_vars = self.num_vars * self.groups
    avg_probs = torch.softmax(x.view(bsz * tsz, self.groups, -1).float(), dim=-1).mean(dim=0)
    quantize_prob_ppl = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-7), dim=-1)).sum()
    quantize_prob_ppl = (num_vars - quantize_prob_ppl) / num_vars
    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)
    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)
    cur_codebook_temp = self.curr_temp
    if not self.time_first:
        x = x.transpose(1, 2)  # BTC -> BCT
    if return_ids:
        hard_x_max = hard_x.argmax(-1).reshape(bsz, tsz, -1)
        # BxTxG
        # create single id from multiple group ids
        target_ids = hard_x.new_zeros(bsz, tsz).long()
        for i in range(self.groups):
            target_ids *= self.num_vars
            target_ids += hard_x_max[:, :, i]
        return x, quantize_prob_ppl, cur_codebook_temp, target_ids
    else:
        return x, quantize_prob_ppl, cur_codebook_temp