Add attention visualization

[gchhablani]

This will depend on whether the model has attention layers on not. Probably need to assume layers and attention.