What is the data type of your index tensor here? You might want to call .long() on it to change the datatype before passing it to gather. (e.g., if it is batch_action, you may want to use batch_action.long().