Freezing when getting out of loop

Hey everyone
I have this code that runs well for most of the inputs but in some cases, in the loop marked with j counter, it freezes at the time when it has to break out of the loop. I follow the memory and CPU usage during the iterations and they’re stable but when the code wants to break of the for loop it freezes and the CPU usage goes form about 50 to 90% and RAM increases from 50 to 80%. I don’t think it’s leakage but it still might be. How should I treat this problem?
Here is the code:

using Plots, Distributions, Statistics, Random, LinearAlgebra, StatsPlots, Pkg, Kronecker, Optim, Interpolations

include(“tauchenHussey.jl”)

function VF_endo_hr(r_l=0167,r=0.2, σ_re=.15, χ_h=0.035, a_num_grid=1500, v0_a=0, v=0, a_l=-1, γ=1.5, β=0.96, ρ=0.94, σ =sqrt(0.04),
a_L = 500, α = 0.3, δ=0.1, y_num_grid= 5, re_num_grid = 5)

#=
v0_a=0; a_l=-1; v=0; r_l = 0.0167; r = 0.2;
a_num_grid=1500; γ=1.5; β=0.96; ρ=0.94; σ =sqrt(0.04);
a_L = 500; α = 0.3; δ=0.1; y_num_grid= 5;a_l=-1;re_num_grid = 5;σ_re=.15; χ_h=0.035;
=#

ϵ_v = 1.0e-6;
ϵ_d = 1.0e-8;

a_grid = range(a_l, stop = a_L, length = a_num_grid); 
#a_grid = exp.(a_grid).-(1-a_l);

w = (α/(r+δ))^(α/(1-α)) *(1-α);



max_iter = 10000;

y,transition = tauchenHussey(y_num_grid,ρ,σ);

y = exp.(y);

dist_perm = transition^1000; ### the stationory distribution of the income shock can be caclculated 
                                ### by multiplying the transition matrix by itself

y = y./(sum(y.*dist_perm[1,:])); ### is reweighted so it's epxection is equal to one

y_1 = kronecker(y, ones(a_num_grid,1));
a_grid_1 = repeat(a_grid, outer = [y_num_grid, 1]); # the grid point for income level

## here I introduce the the iid return grid; 

re_grid = range(-σ_re, stop = σ_re, length = re_num_grid);
re_grid_aux = (re_grid[2:end]+re_grid[1:end-1])./2;

### getting the prob of each grid
d_re =  Normal(0,σ_re);
re_cdf = cdf(d_re, re_grid);
re_cdf_aux = cdf(d_re, re_grid_aux);
re_probs = copy.(re_cdf);
re_probs[1] = re_cdf_aux[1];
re_probs[2:re_num_grid-1] = re_cdf_aux[2:re_num_grid-1].-re_cdf_aux[1:re_num_grid-2];
re_probs[re_num_grid] = 1-re_cdf_aux[re_num_grid-1];

χ = χ_h/(a_L-50);
r_x1 = χ.*(a_grid_1).*(a_grid_1.>0);

if v==0
    if r_l>0.01
        c_temp = a_grid_1.*(r_l.+r_x1)+w.*y_1;
    else
        c_temp = a_grid_1.*(0.01.+r_x1)+w.*y_1;
    end
    c_temp[c_temp.<=0] .= 10000;
    v = (c_temp.^(1.0-γ)./(1.0-γ))./(1.0-β); # frist guess of value function is that consume the income in each period with 
                                            # a fixed income
    v[c_temp.==10000] .= -Inf;
end



a_grid_t = transpose(a_grid);

r_x = χ.*(a_grid_t).*(a_grid_t.>0);

v0 = transpose(reshape(v,a_num_grid,y_num_grid));

#c_expanded  = zeros(y_num_grid,1+a_num_grid);




coh_today = zeros(y_num_grid,a_num_grid);

c =0;
jj = 0;
kk = 0;
accuracy = 0;
#max_iter*10^-2
if v0_a ==0
    v0_a = zeros(y_num_grid,a_num_grid);
    
    v0_a[:,1:a_num_grid-1] .= (v0[:,2:a_num_grid].-v0[:,1:a_num_grid-1])./
            (transpose(a_grid[2:a_num_grid]).-transpose(a_grid[1:a_num_grid-1]));
    v0_a[:,a_num_grid] .= v0_a[:,a_num_grid-1] + 
            (a_grid[a_num_grid]-a_grid[a_num_grid-2])./(a_grid[a_num_grid-1]-a_grid[a_num_grid-2]).*
            (v0_a[:,a_num_grid-1].-v0_a[:,a_num_grid-2]);
    v0_a[isnan.(v0_a)] .= Inf;
    v0_a[v0_a .== Inf] .= 10000;
    
end

ii = 0;
point_add = 0;
tt = 0;
coh_grid = zeros(y_num_grid,a_num_grid,re_num_grid);
re_3d = zeros(1,1,re_num_grid);
re_3d[1,1,:] = re_grid;
re_3d_p = zeros(1,1,re_num_grid);
re_3d_p[1,1,:] .= re_probs;



coh_grid = (1+r_l.+re_3d.+r_x).*a_grid_t.+w.*y;


saving_ong_exp = zeros(y_num_grid,a_num_grid);
saving_ong = zeros(y_num_grid,a_num_grid,re_num_grid);
c_ong = zeros(y_num_grid,a_num_grid,re_num_grid);
v0_temp = zeros(y_num_grid,a_num_grid,re_num_grid);
v0_a_temp = zeros(y_num_grid,a_num_grid,re_num_grid);

for j in 1:max_iter
    

    jj = j;
    println(jj)

    c = (β .* transition*v0_a).^(-1/γ);    ### here we get the consumption decisions
    coh_today = c.+a_grid_t;  ### here I get the COH that has lead to today's consumption and saving





    ### next we have to find the saving and consumption for on-grid points

    for i in 1:y_num_grid
        ii =i;
        coh_grid_t = coh_grid[i,:,:];
        coh_grid_pull = coh_grid_t[:];
        min_coh_temp = minimum(coh_grid_pull);
        min_node = min(min_coh_temp,coh_today[i,1]);
        max_coh_temp = maximum(coh_grid_pull);
        max_node = max(max_coh_temp,coh_today[i,a_num_grid]);
        point_add = a_grid[a_num_grid] + 
                    .5*(max_node-coh_today[i,a_num_grid])./(coh_today[i,a_num_grid]-coh_today[i,a_num_grid-1]).*
                    (a_grid[a_num_grid].-a_grid[a_num_grid-1]); ### here I say that for COHs below the min of COH_today
                    ### we shall save the minimum amount and for the amounts more than max of COH_today we will save 
                    ### linearly more
        nodes = [min_node ;coh_today[i,:] ;max_node];
        Interpolations.deduplicate_knots!(nodes);
    
        itp = interpolate((nodes,),[a_grid[1];a_grid;point_add], Gridded(Linear()));

        saving_pull_temp = itp.(coh_grid_pull); ### now I interpolate the function of COH_today->savings over all 
                                        ### possible COHs on grid which for each y shock is a vector of a_num_grid
                                        ### × re_num_grid points
        saving_ong[i,:,:] .= reshape(saving_pull_temp,a_num_grid,re_num_grid); ### here I reshape the saving on-grid to 
                                        ### to put it in to the matrix format again, after that I can have r_e probs 
                                        ### multiplied to it and get the expected savings (and other variables)
        #=
        if jj == 321
            println(ii)
        end
        =#
    end
        
    c_ong = coh_grid.-saving_ong;

        
    v0_temp .= re_3d_p.*c_ong.^(1-γ)./(1-γ);
    v0_a_temp .= re_3d_p.*(1+r_l.+re_3d.+r_x).*c_ong.^(-γ);

    
        
    saving_is_temp = sum(saving_ong.*re_3d_p,dims=3);
    saving_ong_exp =  saving_is_temp[:,:,1];
        
        
    #end

    
    
    v0_new = sum(v0_temp,dims=3).+(β.*(transition*v0));
    accuracy = maximum(abs.(v0_new-v0));
    if maximum(abs.(v0_new-v0))<ϵ_v
        println("value function converged in $jj with accuaracy $accuracy")
        break;
    end

    v0 = v0_new[:,:,1];

    v0_a_tem = sum(v0_a_temp,dims=3);
    v0_a = v0_a_tem[:,:,1];

    saving_is_temp = nothing;
    v0_a_tem = nothing;
    #println(accuracy)
end

println("value function converged in $jj with accuaracy $accuracy")

saving_tran = transpose(saving_ong_exp);
saving_vector = saving_tran[:];
saving_vector[saving_vector.>a_L] .= a_L-0.00001;
saving_vector[saving_vector.<a_l] .= a_l;
agrid_mat = repeat(a_grid_t,outer=[y_num_grid*a_num_grid,1]);
a_grid_R = transpose([a_grid_t[1];a_grid_t[1:a_num_grid-1]]);
a_grid_L = transpose([a_grid_t[2:a_num_grid];a_grid_t[a_num_grid]]);



policy_r_ind = (saving_vector.<a_grid_L).*(saving_vector.>=a_grid_t); ## here we get indicator of the grid point before the policy point
policy_l_ind = (saving_vector.>a_grid_R).*(saving_vector.<=a_grid_t); ## here we get indicator of the grid point after the policy point
policy_r = sum(policy_r_ind.*agrid_mat,dims=2); # this gives the amount of asset before the policy point
policy_l = sum(policy_l_ind.*agrid_mat,dims=2); # this gives the amount of asset after the policy point

policy = policy_r_ind.*(policy_l.-saving_vector)./(policy_l.-policy_r).+
    policy_l_ind.*(saving_vector.-policy_r)./(policy_l.-policy_r);

dist = 1/( a_num_grid*y_num_grid).*ones(1,a_num_grid*y_num_grid);

policy_rep = repeat(policy, outer = [1, y_num_grid]);
transition_expanded = kronecker(transition,ones(a_num_grid,a_num_grid));
Transition_matrix = policy_rep.*transition_expanded;
dist = 1/( a_num_grid*y_num_grid).*ones(1,a_num_grid*y_num_grid);
jj = 0;
accuracy = 10;
for j in 1:max_iter

    dist_new = dist*Transition_matrix; # the dist is transitioned by the matrix to the new dist
    accuracy =maximum(abs.(dist_new-dist)); 
    if accuracy<ϵ_d

        break;
    end
    dist = dist_new;
    jj =j;

end
println("distribution converged in $jj with accuracy $accuracy")

asset_descion = sum(policy.*a_grid_t,dims=2);
a_dist = asset_descion.*transpose(dist);


a_T = sum(a_dist);
vv = transpose(v0)
v = vv[:];
return v0, dist, asset_descion, a_dist, a_T, policy, a_grid_t, v0_a, saving_vector, v, re_3d, re_3d_p,r_x, saving_ong

end

v0, dist, asset_descion, a_dist, a_T, policy, a_grid_t, v0_a, saving_vector, v, re_3d, re_3d_p,r_x, saving_ong = VF_endo_hr();