A small question on the attn_ada_rmsnorm and attn_ada_layerscale

[Wonder1905]

I hope its the right place to write those question (I think its better than "issues")
In the block of code (for example line 1122):

            action_tokens = attn_ada_rmsnorm(action_tokens, time_cond)
            (state_attn_out, actions_attn_out), (state_keys, state_values, action_keys, action_values) = attn(state_tokens, action_tokens, rotary_emb = rotary_emb, flex_attn_fn = flex_attn_fn, actions_value_residual = actions_value_residual, mask = mask, return_keys_values = True)

            state_cached_keys_values.append((state_keys, state_values))

            actions_value_residual = default(actions_value_residual, action_values)

            action_tokens = attn_ada_layerscale(action_tokens, time_cond)

            state_tokens = state_tokens + state_attn_out
            action_tokens = action_tokens + actions_attn_out`

Does:
action_tokens = attn_ada_layerscale(action_tokens, time_cond)
should be on actions_attn_out or something else then action_tokens?

Thanks a lot!

[lucidrains]

@Wonder1905 ah yes, thank you yet again!