From 9f8bef4fab66bc4d6da31365ef0707143d4d8cdc Mon Sep 17 00:00:00 2001
From: Wei Du <wedu@nvidia.com>
Date: Mon, 30 Jun 2025 12:59:58 -0700
Subject: [PATCH] fix: prevent divisible error by dropping last batch in loader

Signed-off-by: Wei Du <wedu@nvidia.com>
---
 nemo_rl/algorithms/grpo.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py
index ea99de1538..3bcad52849 100644
--- a/nemo_rl/algorithms/grpo.py
+++ b/nemo_rl/algorithms/grpo.py
@@ -191,6 +191,7 @@ def setup(
         batch_size=grpo_config["num_prompts_per_step"],
         shuffle=False,
         collate_fn=rl_collate_fn,
+        drop_last=True,
     )
     if last_checkpoint_path is not None:
         dataloader_state_dict = torch.load(