fix empty tensor (#6319)

Co-authored-by: Tong Li <tong.li35271158@gmail.com>
This commit is contained in:
Tong Li 2025-05-20 17:41:44 +08:00 committed by GitHub
parent bcf2459db5
commit 32afa7bf29
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -254,7 +254,7 @@ class GRPOConsumer(BaseConsumer):
true_indices = torch.nonzero(effective_prompts_mask)
# Make sure the indices are not empty.
if true_indices.numel() > 0:
true_indices = true_indices.squeeze()
true_indices = true_indices.squeeze(-1)
if excessive_prompts_per_rank <= len(true_indices):
excessive_prompts_idx = true_indices[-excessive_prompts_per_rank:]
else: