diff --git a/node.py b/node.py
index 7403099..b456039 100644
--- a/node.py
+++ b/node.py
@@ -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