diff --git a/node.py b/node.py
index 0110f2d..296a521 100644
--- a/node.py
+++ b/node.py
@@ -886,11 +886,8 @@ class GenConsistent:
                                     if pe is not None:
                                         # Extract positional encoding for the saved key positions
                                         # This mimics the approach in attention_processor_characonsist.py
-                                        image_rotary_emb_fg = (
-                                            pe[0][valid_curr_inds] if pe[0].shape[0] >= len(valid_curr_inds) else pe[0],
-                                            pe[1][valid_curr_inds] if pe[1].shape[0] >= len(valid_curr_inds) else pe[1],
-                                        )
-                                        saved_key_fg = flux_math.apply_rope(saved_key_fg, image_rotary_emb_fg)
+                                        image_rotary_emb_fg = pe[:,:,valid_curr_inds]
+                                        saved_key_fg = flux_math.apply_rope1(saved_key_fg, image_rotary_emb_fg)
                                     
                                     # Concatenate the expanded key-value pairs
                                     k = torch.cat([k, saved_key_fg], dim=2)