From 261a79ebdd1f9569dda6f9e06b146127e392e6d3 Mon Sep 17 00:00:00 2001
From: wassname <1103714+wassname@users.noreply.github.com>
Date: Fri, 10 Jul 2026 13:19:54 +0800
Subject: [PATCH] fix: _to_vector forces CPU fp32 -- vjp path returned cuda
 Vectors, crashing U4 step-2 cosine vs cpu reference

Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>
---
 jsteer/jacobian.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/jsteer/jacobian.py b/jsteer/jacobian.py
index 99f0b02..c03f6a5 100644
--- a/jsteer/jacobian.py
+++ b/jsteer/jacobian.py
@@ -82,9 +82,12 @@ def _resolve_layers(layers, n_layers: int) -> list[int]:
 
 def _to_vector(cfg: SteeringConfig, per_layer: dict[int, Tensor]) -> Vector:
     """Wrap unit directions as a steering-lite Vector (mean_diff's layout:
-    stacked["v"] with leading k=1 dim, shared empty)."""
+    stacked["v"] with leading k=1 dim, shared empty). Always CPU fp32 --
+    Jacobian.pullback is CPU-native but the VJP path accumulates on cuda;
+    without the .cpu() the two paths return device-inconsistent Vectors
+    (Claude: found by U4 step-2 crash, pueue 550)."""
     shared = {l: {} for l in per_layer}
-    stacked = {l: {"v": _unit(v.float()).unsqueeze(0)} for l, v in per_layer.items()}
+    stacked = {l: {"v": _unit(v.float().cpu()).unsqueeze(0)} for l, v in per_layer.items()}
     return Vector(cfg, shared, stacked)