def prepare_train_data(self, which, step=None):
        """Prepare train data for both Generator and Discriminator, each samples_list contains k_label batches of data"""
        assert which == 'D' or which == 'G', 'only support for D and G!!'
        real_samples_list = [
            F.one_hot(self.train_data_list[i].random_batch()['target'][:cfg.batch_size],
                      cfg.vocab_size).float().cuda()
            for i in range(cfg.k_label)]
        if which == 'D':
            assert step is not None, 'missing step'
            gen_samples_list = [self.all_gen_samples_list[i][step * cfg.batch_size:(step + 1) * cfg.batch_size]
                                for i in range(cfg.k_label)]  # get a batch from each label
        else:  # 'G'
            gen_samples_list = [
                self.gen.sample(cfg.batch_size, cfg.batch_size, one_hot=True, label_i=i)
                for i in range(cfg.k_label)]

        return real_samples_list, gen_samples_list
Prepare train data for both Generator and Discriminator, each samples_list contains k_label batches of data
instructor/real_data/catgan_instructor.py