scores = np.matmul(Q, K.transpose(0, 2, 1)) / np.sqrt(d_k)