def get_reward_leakgan(self, sentences, rollout_num, dis, current_k):
        """
        get reward via Monte Carlo search for LeakGAN
        :param sentences: size of batch_size * max_seq_len
        :param rollout_num:
        :param dis:
        :param current_k: current training gen

        :return: reward: batch_size * (max_seq_len / step_size)
        """
        with torch.no_grad():
            batch_size = sentences.size(0)
            rewards = torch.zeros([rollout_num * (self.max_seq_len // self.step_size), batch_size]).float()
            if self.gpu:
                rewards = rewards.cuda()
            idx = 0
            for i in range(rollout_num):
                for t in range(self.max_seq_len // self.step_size):
                    given_num = t * self.step_size + 1  # 1, 5, 9, ..
                    samples = self.rollout_mc_search_leakgan(sentences, dis, given_num)
                    out = dis(samples)
                    out = F.softmax(out, dim=-1)
                    reward = out[:, current_k + 1]
                    rewards[idx] = reward
                    idx += 1

        rewards = rewards.view(batch_size, self.max_seq_len // self.step_size, rollout_num)
        rewards = torch.mean(rewards, dim=-1)
        return rewards

        get reward via Monte Carlo search for LeakGAN
        :param sentences: size of batch_size * max_seq_len
        :param rollout_num:
        :param dis:
        :param current_k: current training gen

        :return: reward: batch_size * (max_seq_len / step_size)
        
utils/rollout.py