Fix multiple subject
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user