Skip to content

Commit b57405c

Browse files
author
Ronald Tse
committed
fix(ride): teacher resident on cuda for displacement - labeling offloads it to cpu
1 parent 61d8a81 commit b57405c

1 file changed

Lines changed: 4 additions & 0 deletions

File tree

‎src/gpu/modal_distill.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1246,6 +1246,10 @@ def _usable(ck: Path) -> bool:
12461246
print(f"[{spec_id}] resume training from step-{step}", flush=True)
12471247

12481248
if ride is not None:
1249+
# displacement needs the teacher resident on cuda for both the
1250+
# ridge fit and every training step (labeling offloads it)
1251+
teacher.to("cuda")
1252+
teacher.eval()
12491253
# Ridge map per selected layer: teacher hidden space -> the
12501254
# student's INITIAL hidden space (linearity preserves the
12511255
# displacement through the projection). Fit on the first

0 commit comments

Comments
 (0)