diff --git a/trlx/trainer/accelerate_ppo_trainer.py b/trlx/trainer/accelerate_ppo_trainer.py index f004e054..424d3d87 100644 --- a/trlx/trainer/accelerate_ppo_trainer.py +++ b/trlx/trainer/accelerate_ppo_trainer.py @@ -352,9 +352,9 @@ def make_experience(self, num_rollouts: int = 1024, iter_count: int = 0): # noq scores_mask = scores != -np.inf if self.config.train.reward_only_in_main_process: - str_samples, str_prompts, str_outputs = self.decode(prompt_tensors, samples, append_eos_token=True) + _, _, str_outputs = self.decode(prompt_tensors, samples, append_eos_token=True) else: - str_samples, str_prompts, str_outputs = all_str_samples, all_str_prompts, all_str_outputs + str_outputs = all_str_outputs # Pad the sample outputs outputs = self.tokenizer(str_outputs).input_ids