MPS Entanglement Entropy

Hi, I am quite new to using ITensors and Julia in general, so my question may sound a bit naive. I am sorry about this. I am trying to write a code that computes the entanglement entropy of an MPS for a certain bond b before and after the application of a two-sites gate (on the sites linked by b of course). Thanks to the documentation and to other posts, I managed to write the following function:

# Function to compute the entanglement entropy for a given MPS on bond b
function MPS_ent_entropy(psi,b)
    if b == 1
        psi =  orthogonalize!(psi, b)
        _,S = svd(psi[b], (siteind(psi,b),))
        SvN = 0.0
        for n=1:dim(S, 1)
          p = S[n,n]^2
          SvN -= p * log(p)    
        end    
        return SvN
    else
        psi =  orthogonalize!(psi, b)
        _,S = svd(psi[b], (linkind(psi, b-1), siteind(psi,b)))
        SvN = 0.0
        for n=1:dim(S, 1)
          p = S[n,n]^2
          SvN -= p * log(p)    
        end
        return SvN
    end
end    

but then if I generate a random product MPS phi and apply a CX gate on two of its sites, I get a NaN entanglement entropy.

N = 10
sites = siteinds("S=1/2", N)
phi = productMPS(sites, "Up")
CX = [1 0 0 0 ; 0 1 0 0 ; 0 0 0 1 ; 0 0 1 0]
CX = op("CX",sites[3],sites[4]) # create a 2-qubits tensor
phi2 = apply(CX,phi)
MPS_ent_entropy(phi2,3)

Can someone help me in figuring out where the problem is?

Your p value is exactly zero and causing the NaN, you should put in an if statement to only add values p>0