The split heads should be concatenated in dim=2 Former-commit-id: 4b3f247f270d44df9fe226cfe0dabfb7fcd2deda
Former-commit-id: e1a8431770fc36c0c9ee7fed4abbc3d7fdcc5efd