Function split_heads 

pub fn split_heads(
    qkv: &Tensor,
    n_head: usize,
) -> Result<(Tensor, Tensor, Tensor)>

qkv [t, 3c] -> (q, k, v) each [h, t, c/h].