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

for more information, see https://pre-commit.ci

studio/backend/core/export/export.py

@@ -120,7 +120,8 @@ def _is_oom_error(exc: BaseException) -> bool:
     """
     if torch is not None:
         oom_types = tuple(
-            t for t in (
+            t
+            for t in (
                 getattr(torch, "OutOfMemoryError", None),
                 getattr(getattr(torch, "cuda", None), "OutOfMemoryError", None),
                 getattr(getattr(torch, "xpu", None), "OutOfMemoryError", None),
@@ -502,11 +503,7 @@ class ExportBackend:
             # GPU, so when a training or chat job already owns the others (which
             # routes/export.py deliberately allows) the shard can OOM where the pre-#7053
             # single-device load succeeded. Fall back once before giving up.
-            if (
-                _device_map_override is None
-                and _is_oom_error(e)
-                and _multi_gpu_device_map_kwargs()
-            ):
+            if _device_map_override is None and _is_oom_error(e) and _multi_gpu_device_map_kwargs():
                 # Retry outside this block: the live traceback pins the half-built
                 # model's frames, so an in-block retry inherits the exhausted device.
                 oom_retry_reason = str(e)

tests/test_compressed_export_gpu_release.py

@@ -43,13 +43,12 @@ class _FakeLogger:
 def _load_helpers(fake_torch, fake_logger):
     tree = ast.parse(_SAVE_PY.read_text(encoding = "utf-8"))
     keep = [
-        node for node in tree.body
+        node
+        for node in tree.body
         if (isinstance(node, ast.FunctionDef) and node.name in _WANTED)
         or (
             isinstance(node, ast.Assign)
-            and any(
-                isinstance(t, ast.Name) and t.id in _WANTED_ASSIGNS for t in node.targets
-            )
+            and any(isinstance(t, ast.Name) and t.id in _WANTED_ASSIGNS for t in node.targets)
         )
     ]
     n_fns = sum(1 for node in keep if isinstance(node, ast.FunctionDef))
@@ -101,9 +100,7 @@ def _fake_accelerate(monkeypatch):
     accel.dispatch_model = _dispatch
     hooks = types.ModuleType("accelerate.hooks")
     hooks.remove_hook_from_submodules = lambda model: calls["removed"].append(model)
-    hooks.add_hook_to_module = lambda module, hook: calls["hooks_added"].append(
-        (module, hook)
-    )
+    hooks.add_hook_to_module = lambda module, hook: calls["hooks_added"].append((module, hook))
     accel.hooks = hooks
     monkeypatch.setitem(sys.modules, "accelerate", accel)
     monkeypatch.setitem(sys.modules, "accelerate.hooks", hooks)
@@ -341,7 +338,11 @@ def test_torchao_export_uses_the_shared_release():
 class _Child:
     """Minimal stand-in for an nn.Module leaf, enough for the dispatch walk."""
 
-    def __init__(self, name = "inner", device_map = None):
+    def __init__(
+        self,
+        name = "inner",
+        device_map = None,
+    ):
         self._modules = {}
         self.__dict__["_name"] = name
         if device_map is not None:
@@ -449,11 +450,11 @@ def test_snapshot_restores_a_forward_patched_after_the_dispatch(_fake_accelerate
     mlp = _Child(name = "mlp")
     root._modules["mlp"] = mlp
 
-    stock_forward = lambda *a, **k: "stock"          # noqa: E731
+    stock_forward = lambda *a, **k: "stock"  # noqa: E731
     fused_forward = lambda *a, **k: "unsloth-fused"  # noqa: E731
     mlp._hf_hook = object()
-    mlp._old_forward = stock_forward   # captured by accelerate at dispatch time
-    mlp.forward = fused_forward        # installed by unsloth afterwards
+    mlp._old_forward = stock_forward  # captured by accelerate at dispatch time
+    mlp.forward = fused_forward  # installed by unsloth afterwards
 
     snapshot = ns["_snapshot_dispatch_state"](root)
 

unsloth/save.py

@@ -4638,7 +4638,8 @@ def _accelerate_dispatch_root(model):
         children = getattr(node, "__dict__", {}).get("_modules") or {}
         nxt = next(
             (
-                children[a] for a in ("base_model", "model")
+                children[a]
+                for a in ("base_model", "model")
                 if hasattr(children.get(a), "named_modules")
             ),
             None,
@@ -4836,6 +4837,7 @@ def _restore_model_after_quantize_subprocess(model, restore_token) -> None:
                 _restore_dispatch_state(root, snapshot)
             else:
                 from accelerate import dispatch_model
+
                 # skip_keys matters: without it accelerate moves every forward kwarg
                 # to the executing device, wrong for device-invariant cache tensors.
                 dispatch_model(