ilkarman:

could the indexing be creating a break in autograd?


It should’t break the computation graph. Applying softmax and then indexing the weights is a valid sequence of ops. Additional code would be helpful (the whole forward pass and training loop) to solve the issue.