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 sure that you don’t share memory between the version used by gather and the one used in the rest of your code. That way inplace operations in the rest of your code won’t be a problem for gather.