[BUG] Error when wrapping the step

[sdiazlor]

Describe the bug
A clear and concise description of what the bug is.
When working on this notebook (#949), I encountered that with the dev branch, the task wasn't working.

Maybe related to this one too: #991

To Reproduce
Code to reproduce

labels_topic= ["World", "Sports", "Sci/Tech", "Business"]
labels_fact_opinion = ["Fact-based", "Opinion-based"]

task_templates = [
    "Determine the news article as {}",
    "Classify news article as {}",
    "Identify the news article as {}",
    "Categorize the news article as {}",
    "Label the news article using {}",
    "Annotate the news article based on {}",
    "Determine the theme of a news article from {}",
    "Recognize the topic of the news article as {}",
]

classification_tasks = [
    {"task": action.format(" or ".join(random.sample(labels_topic, 2)))}
    for action in task_templates for _ in range(4)
] + [
    {"task": action.format(" or ".join(random.sample(labels_fact_opinion, 2)))}
    for action in task_templates
]

difficulties = ["college", "high school", "PhD"]
clarity = ["clear", "understandable with some effort", "ambiguous"]

with Pipeline("texcat-generation-pipeline") as pipeline:

    tasks_generator = LoadDataFromDicts(data=classification_tasks)

    generate_data = []
    for difficulty in difficulties:
        for clarity_level in clarity:
            task = GenerateTextClassificationData(
                language="English",
                difficulty=difficulty,
                clarity=clarity_level,
                num_generations=2,
                llm=InferenceEndpointsLLM(
                    model_id="meta-llama/Meta-Llama-3.1-8B-Instruct",
                    tokenizer_id="meta-llama/Meta-Llama-3.1-8B-Instruct",
                    generation_kwargs={"max_new_tokens": 512, "temperature": 0.7},
                ),
                input_batch_size=5,
            )
            generate_data.append(task)

    for task in generate_data:
        tasks_generator.connect(task)

Expected behaviour
A clear and concise description of what you expected to happen.

Screenshots
If applicable, add screenshots to help explain your problem.

Desktop (please complete the following information):

• Package version: 1.4dev
• Python version: 3.11.4

Additional context
Add any other context about the problem here.
Note: Already shared on Slack, just realized I didn't create an issue

[plaguss]

Hi @sdiazlor, can you share the classification_tasks data to run the example?

[sdiazlor]

@plaguss thanks, I updated the example code. With the 1.3.2 works well.

[plaguss]

Can you test with the fix and see if it works?

[sdiazlor]

It works now, thanks for tackling this!