fix

Signed-off-by: Hemil Desai <hemild@nvidia.com>

Author: hemildesai

nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py

@@ -661,6 +661,10 @@ def score(self, data: BatchedDataDict) -> BatchedDataDict[ScoreOutputSpec]:
                         sequence_dim=sequence_dim,
                     )
 
+                # skip keeping the scores for the dummy batches
+                if batch_idx >= iterator_len:
+                    continue
+
                 all_rm_scores.append(rm_scores)
 
         all_rm_scores = torch.cat(all_rm_scores, dim=0)
@@ -746,6 +750,10 @@ def get_topk_logits(
                         sequence_dim=sequence_dim,
                     )
 
+                # skip keeping the topk values for the dummy batches
+                if batch_idx >= iterator_len:
+                    continue
+
                 # Keep only real sequence tokens (no trimming here; padded positions can be masked downstream)
                 # Shapes remain [B, S, k].
                 out_topk_vals.append(vals.cpu())