From 417688974d5cd3634580dbea9b470f97c3186bbc Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Tue, 1 Sep 2026 06:37:47 +0200 Subject: [PATCH] fix(e5): move the mtp head to cuda with the student MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit build_mtp ran while the student was still on cpu (student.to('cuda') happens after label generation), so the head's Linears stayed on cpu and the first training step died in a device mismatch. Labels survive in the run dir — the relaunch skips labeling and trains. --- src/gpu/modal_distill.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/gpu/modal_distill.py b/src/gpu/modal_distill.py index 6000018..7ba823d 100644 --- a/src/gpu/modal_distill.py +++ b/src/gpu/modal_distill.py @@ -1033,6 +1033,8 @@ def accept_label(src: str, label: str) -> None: torch.cuda.empty_cache() student.to("cuda") student.gradient_checkpointing_enable() + if mtp_head is not None: # built while student was still on cpu + mtp_head.to("cuda") class TeacherPairs(Dataset): def __len__(self):