fix: taste direction works on the CPU, so Taste prompt nudge / score slider no longer die on a cuda-vs-cpu mismatch

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
2026-10-17 01:11:50 +03:00
co-authored by Claude Sonnet 5.5
parent 388ca59428
commit 03d902e208
+3 -2
View File
@@ -11,13 +11,14 @@ MIN_LIKED = 3 # v4: direction_count >= 3
def direction(rows, current=None, similar=False):
"""-> (unit vector | None, how it was found)."""
pooled = [(r["rows"][NAME].float(), r["reward"]) for r in rows if NAME in r["rows"]]
pooled = [(r["rows"][NAME].detach().float().cpu(), r["reward"]) for r in rows if NAME in r["rows"]]
liked = [p for p, w in pooled if w > 0]
if len(liked) < MIN_LIKED:
return None, f"needs {MIN_LIKED} liked clips ({len(liked)} so far)"
mean = torch.stack([p for p, _w in pooled]).mean(0)
if similar and current is not None:
near = sorted(((float(F.cosine_similarity(p, current.float(), dim=0)), p) for p in liked),
current = current.detach().float().cpu() # the stored rows live on the CPU; the live prompt is on the GPU
near = sorted(((float(F.cosine_similarity(p, current, dim=0)), p) for p in liked),
key=lambda s: -s[0])[:NEAREST]
near = [(s, p) for s, p in near if s >= MIN_SIM]
if near: