Keep attention profiling compatible with torch 1.x (#159)