The application of self-attention layer raised index error

Question

So I am doing a classification machine learning with the input of (batch, step, features).

In order to improve the accuracy of this model, I intended to apply a self-