gemini

Author: hjh0119

swift/trainers/rlhf_trainer/grpo_trainer.py

@@ -2062,10 +2062,7 @@ def get_chunked_inputs(self, inputs, start_idx, end_idx):
         # for LLM, slice the inputs
         for key, val in inputs.items():
             if isinstance(val, torch.Tensor):
-                if val.ndim == 0:
-                    chunk_inputs[key] = val
-                else:
-                    chunk_inputs[key] = val[start_idx:end_idx]
+                chunk_inputs[key] = val if val.ndim == 0 else val[start_idx:end_idx]
             else:
                 chunk_inputs[key] = val
         if self.is_multimodal: