Fix multiple subject

This commit is contained in:
ZHOU He
2025-11-17 22:40:00 +08:00
parent 7fe88b0b07
commit 6790ea095d
+9 -9
View File
@@ -729,14 +729,14 @@ class GetCrossSim:
argmax_indices[0],
max_sim[0]>sim_threshold
)
all_id_fg_inds.append(id_fg_inds)
all_curr_fg_inds.append(curr_fg_inds)
all_id_fg_inds.append(id_fg_inds.cpu())
all_curr_fg_inds.append(curr_fg_inds.cpu())
all_max_sim.append(max_sim)
all_argmax_indices.append(argmax_indices)
all_cross_sim.append(avg_sim)
cross_sim["id_fg_inds"] = torch.stack(all_id_fg_inds).cpu()
cross_sim["curr_fg_inds"] = torch.stack(all_curr_fg_inds).cpu()
cross_sim["id_fg_inds"] = all_id_fg_inds
cross_sim["curr_fg_inds"] = all_curr_fg_inds
cross_sim["max_sim"] = torch.cat(all_max_sim).cpu()
cross_sim["argmax_indices"] = torch.cat(all_argmax_indices).cpu()
cross_sim["debug_mask"] = torch.cat(all_cross_sim).cpu()
@@ -863,8 +863,8 @@ class GenConsistent:
# Process all subjects, not just the last one (index = len(subject) - 1)
for subj_idx in range(len(subject)):
# Check if we have valid cross-sim data and subject data to work with
if (subj_idx in id_fg_inds and len(id_fg_inds[subj_idx]) > 0 and
subj_idx in curr_fg_inds and len(curr_fg_inds[subj_idx]) > 0 and
if (len(id_fg_inds[subj_idx]) > 0 and
len(curr_fg_inds[subj_idx]) > 0 and
sampling_step in subject[subj_idx] and
layer in subject[subj_idx][sampling_step]["layers"] and
"k" in subject[subj_idx][sampling_step]["layers"][layer] and
@@ -923,9 +923,9 @@ class GenConsistent:
saved_x = subject[subj_idx][sampling_step]["layers"][layer]["x"].to(x.device, dtype=attn_dtype)
# Ensure we have valid indices within tensor bounds
if subj_idx < id_fg_inds.shape[0] and subj_idx < curr_fg_inds.shape[0] and subj_idx < max_sim.shape[0]:
subj_id_fg_inds = id_fg_inds[subj_idx] if id_fg_inds.dim() > 1 else id_fg_inds
subj_curr_fg_inds = curr_fg_inds[subj_idx] if curr_fg_inds.dim() > 1 else curr_fg_inds
if subj_idx < len(id_fg_inds) and subj_idx < len(curr_fg_inds) and subj_idx < max_sim.shape[0]:
subj_id_fg_inds = id_fg_inds[subj_idx]
subj_curr_fg_inds = curr_fg_inds[subj_idx]
subj_max_sim = max_sim[subj_idx]
# Validate indices