fix: formatting

Signed-off-by: Abhishek <[email protected]>
huggingface · Oct 3, 2024 · a27b3a2 · a27b3a2
1 parent 56e32cb
commit a27b3a2
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/trl/trainer/utils.py b/trl/trainer/utils.py
@@ -233,11 +233,15 @@ def torch_call(self, examples: List[Union[List[int], Any, Dict[str, Any]]]) -> D
             batch["labels"][batch["position_ids"] == 0] = self.ignore_index
 
             flattened_position_ids = batch["position_ids"].flatten()
-            indices_q = torch.arange(flattened_position_ids.size(0), device=flattened_position_ids.device, dtype=torch.int32)
+            indices_q = torch.arange(
+                flattened_position_ids.size(0), device=flattened_position_ids.device, dtype=torch.int32
+            )
             batch["cu_seq_lens_q"] = torch.cat(
                 (
                     indices_q[flattened_position_ids == 0],
-                    torch.tensor(flattened_position_ids.size(), device=flattened_position_ids.device, dtype=torch.int32),
+                    torch.tensor(
+                        flattened_position_ids.size(), device=flattened_position_ids.device, dtype=torch.int32
+                    ),
                 )
             )
             batch["cu_seq_lens_k"] = batch["cu_seq_lens_q"]