Thank you very much, Tony-Y.
The index_select function cannot be diff. so the gradient cannot backprop. to the previous S_K architecture.
My problem is how do I implement  “select” action in pytorch instead of using the index_select function to implement it?