def step(self, inp, hidden, label=None):
        """
        RelGAN step forward
        :param inp: batch_size
        :param hidden: memory size
        :param label: batch_size, specific label index
        :return: pred, hidden, next_token
            - pred: batch_size * vocab_size, use for adversarial training backward
            - hidden: next hidden
            - next_token: [batch_size], next sentence token
        """
        assert type(label) == torch.Tensor, 'missing label'
        emb = self.embeddings(inp).unsqueeze(1)

        # cat category vector
        label_onehot = F.one_hot(label, self.k_label).float()  # batch_size * k_label
        label_onehot_ex = label_onehot.unsqueeze(1).expand(-1, 1, -1)  # batch_size * 1 * k_label
        label_vec = torch.bmm(label_onehot_ex, self.cat_mat.expand(inp.size(0), -1, -1))  # batch_size * 1 * embed_dim
        emb = torch.cat((emb, label_vec), dim=-1)  # batch_sie * len * (k_label + embed_dim)

        out, hidden = self.lstm(emb, hidden)
        gumbel_t = self.add_gumbel(self.lstm2out(out.squeeze(1)))
        next_token = torch.argmax(gumbel_t, dim=1).detach()

        pred = F.softmax(gumbel_t * self.temperature, dim=-1)  # batch_size * vocab_size

        return pred, hidden, next_token

        RelGAN step forward
        :param inp: batch_size
        :param hidden: memory size
        :param label: batch_size, specific label index
        :return: pred, hidden, next_token
            - pred: batch_size * vocab_size, use for adversarial training backward
            - hidden: next hidden
            - next_token: [batch_size], next sentence token
        
models/CatGAN_G.py