PyTorch Self-Attention

Manually implement the Self-Attention mechanism in PyTorch. Specify the dimensions of several important matrices and compute the time complexity. When applying multiple layers of Self-Attention, how do you stabilize the training of Transformers and solve the issues of gradient explosion and GPU memory explosion? Example Input (2, 10, 64) [(batch_size, seq_len, d_k)]