From 1b5de29e71a581db47fcc6824cf48b5ef67bff36 Mon Sep 17 00:00:00 2001
From: Adria Blancafort <76774853+adriablancafort@users.noreply.github.com>
Date: Wed, 7 Jan 2026 18:08:57 +0100
Subject: [PATCH] Fix undefined variable in chat_rl after recent refactor

* Fix undefined variable

* Remove unused import

Remove unused import 're' from chat_rl.py
---
 scripts/chat_rl.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/scripts/chat_rl.py b/scripts/chat_rl.py
index 1a09962..ad557b9 100644
--- a/scripts/chat_rl.py
+++ b/scripts/chat_rl.py
@@ -19,7 +19,6 @@ torchrun --standalone --nproc_per_node=8 -m scripts.chat_rl -- --run=default
 import argparse
 import os
 import itertools
-import re
 import wandb
 import torch
 import torch.distributed as dist
@@ -174,7 +173,7 @@ def run_gsm8k_eval(task, tokenizer, engine,
         tokens = tokenizer.render_for_completion(conversation)
         prefix_length = len(tokens)
         # Generate k samples using batched generation inside the Engine
-        assert num_samples <= device_batch_size # usually this is true. we can add a loop if not...
+        assert num_samples <= args.device_batch_size # usually this is true. we can add a loop if not...
         generated_token_sequences, masks = engine.generate_batch(
             tokens,
             num_samples=num_samples,