chore: Fix formatting

Author: bradhilton

src/art/unsloth/train.py

@@ -114,9 +114,9 @@ def compute_loss(
         next_input_ids = shift_tensor(inputs["tokens"], 0)
         chunk_size = _config.get("logprob_calculation_chunk_size", 1024)
         # Assert that sequence length is evenly divisible by the chunk size
-        assert (
-            seq_len % chunk_size == 0
-        ), f"Sequence length ({seq_len}) must be evenly divisible by chunk size ({chunk_size})"
+        assert seq_len % chunk_size == 0, (
+            f"Sequence length ({seq_len}) must be evenly divisible by chunk size ({chunk_size})"
+        )
         os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "1"
         forward_kwargs = {}
         if "pixel_values" in inputs:
@@ -371,9 +371,7 @@ def _calculate_logprobs(
         chunk_logits = torch.matmul(chunk_hs, lm_head_t)  # [B, chunk_size, V]
         chunk_selected_logits = torch.gather(
             chunk_logits, dim=-1, index=chunk_input_ids.unsqueeze(-1)
-        ).squeeze(
-            -1
-        )  # [B, chunk_size]
+        ).squeeze(-1)  # [B, chunk_size]
         chunk_logsumexp = torch.logsumexp(chunk_logits, dim=-1)  # [B, chunk_size]
         log_probs[:, i : i + chunk_size] = chunk_selected_logits - chunk_logsumexp