mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-06-23 06:00:44 +00:00
fix empty tensor (#6319)
Co-authored-by: Tong Li <tong.li35271158@gmail.com>
This commit is contained in:
parent
bcf2459db5
commit
32afa7bf29
@ -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:
|
||||
|
Loading…
Reference in New Issue
Block a user