diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index ff9bc1ec78..ea7183fc5a 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -215,13 +215,19 @@ def grpo_trainer__generate_and_score_completions(function_name, function): # The new multi-line string that will replace the line above replacement_lines = """ batch_size = self.args.per_device_train_batch_size if mode == "train" else self.args.per_device_eval_batch_size - if not has_images: - # Left pad prompt before calculation old and ref hidden states - prompt_completion_ids = left_pack_padding(prompt_completion_ids, self.processing_class.pad_token_id) - self.model.for_training()""" + try: + #TRL 0.23.1 and below path + if not has_images: + # Left pad prompt before calculation old and ref hidden states + prompt_completion_ids = left_pack_padding(prompt_completion_ids, self.processing_class.pad_token_id) + self.model.for_training() + except: + #TRL 0.24.0 and below path + if images is None: + # Left pad prompt before calculation old and ref hidden states + prompt_completion_ids = left_pack_padding(prompt_completion_ids, self.processing_class.pad_token_id) + self.model.for_training()""" - if "has_images" not in function: - raise NotImplementedError("Unsloth: For now we support `trl<=0.23.1`. Please downgrade!") function = function.replace(line_to_replace, replacement_lines) pattern_to_find = re.compile(