From 11dd4710cd7ffa377e47a348310021451ae7811a Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Sun, 30 Aug 2026 18:47:32 +0200 Subject: [PATCH] fix: stitch padding on dimension growth (d768 heads 6->12) --- src/gpu/modal_distill.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/gpu/modal_distill.py b/src/gpu/modal_distill.py index 22c0c88..cc2c2ce 100644 --- a/src/gpu/modal_distill.py +++ b/src/gpu/modal_distill.py @@ -112,9 +112,10 @@ def svd_stitch_state(wide: dict, narrow: dict) -> dict: w = u[:, :o2].T @ w if i2 < i: w = w @ vh[:i2, :].T - if o2 > o or i2 > i: + wo, wi = w.shape + if wo < o2 or wi < i2: padded = torch.zeros((o2, i2), dtype=w.dtype) - padded[: min(o, o2), : min(i, i2)] = w[: min(o, o2), : min(i, i2)] + padded[: min(wo, o2), : min(wi, i2)] = w[: min(wo, o2), : min(wi, i2)] w = padded out[name] = w.to(tgt.dtype) elif src.dim() == 1: