Created by: ktlichkid
Fixed simple_attention implementation to match the implementation in other frameworks.