ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2026-04-21 06:19:54 +00:00
commit a94cece8f6
1 changed files with 6 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

8
scripts/benchmarks/unsloth_grpo_common.py
View file

@ -322,7 +322,9 @@ def write_stats(
json.dump(stats, f, indent = 2)
def maybe_compile_trainer_forwards(trainer, compile_mode, *, dynamic: bool = True, tag: str = ""):
def maybe_compile_trainer_forwards(
trainer, compile_mode, *, dynamic: bool = True, tag: str = ""
):
"""torch.compile wrap `trainer.model.forward` and (if present)
`trainer.ref_model.forward`. No-op if `compile_mode` is falsy.
@ -339,7 +341,9 @@ def maybe_compile_trainer_forwards(trainer, compile_mode, *, dynamic: bool = Tru
except AttributeError:
pass
prefix = f"[{tag}] " if tag else ""
print(f"{prefix}Compiling trainer.model.forward (mode={compile_mode}, dynamic={dynamic})")
print(
f"{prefix}Compiling trainer.model.forward (mode={compile_mode}, dynamic={dynamic})"
)
trainer.model.forward = torch.compile(
trainer.model.forward,
mode = compile_mode,