fix empty tensor (#6319)

Co-authored-by: Tong Li <[email protected]>

Author: TongLi3701

applications/ColossalChat/coati/distributed/grpo_consumer.py

@@ -254,7 +254,7 @@ def step(self, step_idx: int, pbar: Any, **kwargs) -> Optional[float]:
                     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: