Created by: ktlichkid

Fix implementation of simple_attention to match the implementation of other frameworks.