Update PTST.py

PTST.py

@@ -214,7 +214,7 @@ class MultiHeadSelfAttention(nn.Module):
             # S_{xx'} = sum_i R'_{ix} R'_{ix'}
             S  = jnp.einsum("btx,bty->bxy", out, out) / jnp.maximum(T, 1)  # (B,x,x)
             S2 = jnp.einsum("bii->bi", S)
-            S  = S2[None,:] + S2[:,None] - 2*S
+            S  = S2[:,None,:] + S2[:,:,None] - 2*S
             S  = -1.*S
 
             # column-wise softmax (columns sum to 1): softmax over row axis