Fix RuntimeError: pad attn scores back to original query sequence length, instead of unpadded sequence length (i.e. no change).

Correct the output dtype of rmsnorm_func (#13)

remove torch.jit

init