Purifying MPO into MPS

I’m trying to take an MPO representing a density operator and turn it into a doubled MPS via contracting Bell states to the unprimed indexes.

However, when I try to split the newly contracted tensors into two using SVD to get “skinny” indexes some are the right dimension (aka dim 2 for qubits) but others are dim 4.

Here is the code snippet I’m using to do this: I really appreciate the help!

function purify_MPO_to_MPS(M::MPO)::MPS
    N = length(M)
    
    s_unprimed_vec = [siteind(M, j, plev=0) for j in 1:N]
    s_primed_vec = [siteind(M, j, plev=1) for j in 1:N]

   
    doubled_sites = [addtags(s_unprimed_vec[j], "doubled") for j in 1:N]

    doubled_mps_tensors = ITensor[]
    
    for j in 1:N
        mpo_tensor = M[j]
        s_unprimed = s_unprimed_vec[j]
        s_primed = s_primed_vec[j]
        s_doubled_new = doubled_sites[j]

       
        bell_state = ITensor(s_doubled_new, s_primed)
        for k in 1:dim(s_doubled_new)
            bell_state[s_doubled_new(k), s_primed(k)] = 1.0
        end

        
        combined_tensor = mpo_tensor * bell_state
        
        
         U, S, V = svd(combined_tensor, s_unprimed, s_doubled_new)
         lefttensor=U
         righttensor=S*V
        
        push!(doubled_mps_tensors, lefttensor)
        push!(doubled_mps_tensors, righttensor)
    end

    
    purified_psi = MPS(doubled_mps_tensors)
    
    return normalize(purified_psi)
end