Solved by albanD in post #4 
              
              
                Versions tracks how many times it has been changed inplace. 
Inplace modifications are anything like positions[foo] =  or positions.add_(foo) or positions += 2. 
a simple fix you can use is replace your gather call to: 
states = hidden_states.gather(-2, positions.clone())

The clone here will make s…