I face the same problem, I also intend to use pytorch to implement the Kalman filter. After making a lot of attempts, For me, I found this to be the better way than looping：
    # Combine tensors on cpu
    covariances = []
    for row, track_idx in enumerate(track_indices):
        track = tracks[track_idx]
        covariances.append(track.covariance.cpu())

    # Then copy to gpu
    covariances = torch.cat(covariances, dim=0).to("cuda:0")

    # The calculation takes place on the GPU
    torch.cholesky(covariance)