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):