rotary embed in loop

Author: DavidLandup0

keras_hub/src/models/smollm3/smollm3_backbone.py

@@ -118,16 +118,16 @@ def __init__(
         )
 
         cache_update_index = kwargs.get('self_attention_cache_index')
-        print(cache_update_index)
 
         start_index = (
             cache_update_index if cache_update_index is not None else 0
         )
 
         hidden_states = self.token_embedding(token_id_input)
-        position_embeddings = self.rotary_embedding(hidden_states, start_index=start_index)
+        
 
         for decoder_layer in self.transformer_layers[:num_layers]:
+            position_embeddings = self.rotary_embedding(hidden_states, start_index=start_index)
             hidden_states = decoder_layer(
                 hidden_states,
                 position_embeddings=position_embeddings,