diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl
index 28093b54..b537f83a 100644
--- a/src/ext/AMDGPU/2D.jl
+++ b/src/ext/AMDGPU/2D.jl
@@ -85,6 +85,55 @@ function JR2D.PTThermalCoeffs(
     return PTThermalCoeffs(rheology, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL)
 end
 
+function JR2D.update_thermal_coeffs!(
+    pt_thermal::JustRelax.PTThermalCoeffs{T,<:ROCArray}, rheology, phase_ratios, args, dt
+) where {T}
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ,
+        pt_thermal.dτ_ρ,
+        rheology,
+        phase_ratios.center,
+        args,
+        pt_thermal.max_lxyz,
+        pt_thermal.Vpdτ,
+        inv(dt),
+    )
+    return nothing
+end
+
+function JR2D.update_thermal_coeffs!(
+    pt_thermal::JustRelax.PTThermalCoeffs{T,<:ROCArray}, rheology, args, dt
+) where {T}
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ,
+        pt_thermal.dτ_ρ,
+        rheology,
+        args,
+        pt_thermal.max_lxyz,
+        pt_thermal.Vpdτ,
+        inv(dt),
+    )
+    return nothing
+end
+
+function JR2D.update_thermal_coeffs!(
+    pt_thermal::JustRelax.PTThermalCoeffs{T,<:ROCArray}, rheology, ::Nothing, args, dt
+) where {T}
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ,
+        pt_thermal.dτ_ρ,
+        rheology,
+        args,
+        pt_thermal.max_lxyz,
+        pt_thermal.Vpdτ,
+        inv(dt),
+    )
+    return nothing
+end
+
 # Boundary conditions
 function JR2D.flow_bcs!(
     ::AMDGPUBackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions
diff --git a/src/ext/AMDGPU/3D.jl b/src/ext/AMDGPU/3D.jl
index 112347d2..1b45cac4 100644
--- a/src/ext/AMDGPU/3D.jl
+++ b/src/ext/AMDGPU/3D.jl
@@ -85,6 +85,55 @@ function JR3D.PTThermalCoeffs(
     return PTThermalCoeffs(rheology, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL)
 end
 
+function JR3D.update_thermal_coeffs!(
+    pt_thermal::JustRelax.PTThermalCoeffs{T,<:ROCArray}, rheology, phase_ratios, args, dt
+) where {T}
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ,
+        pt_thermal.dτ_ρ,
+        rheology,
+        phase_ratios.center,
+        args,
+        pt_thermal.max_lxyz,
+        pt_thermal.Vpdτ,
+        inv(dt),
+    )
+    return nothing
+end
+
+function JR3D.update_thermal_coeffs!(
+    pt_thermal::JustRelax.PTThermalCoeffs{T,<:ROCArray}, rheology, args, dt
+) where {T}
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ,
+        pt_thermal.dτ_ρ,
+        rheology,
+        args,
+        pt_thermal.max_lxyz,
+        pt_thermal.Vpdτ,
+        inv(dt),
+    )
+    return nothing
+end
+
+function JR3D.update_thermal_coeffs!(
+    pt_thermal::JustRelax.PTThermalCoeffs{T,<:ROCArray}, rheology, ::Nothing, args, dt
+) where {T}
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ,
+        pt_thermal.dτ_ρ,
+        rheology,
+        args,
+        pt_thermal.max_lxyz,
+        pt_thermal.Vpdτ,
+        inv(dt),
+    )
+    return nothing
+end
+
 # Boundary conditions
 function JR3D.flow_bcs!(
     ::AMDGPUBackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions
diff --git a/src/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl
index cbe74022..dfb3aaff 100644
--- a/src/ext/CUDA/2D.jl
+++ b/src/ext/CUDA/2D.jl
@@ -84,6 +84,55 @@ function JR2D.PTThermalCoeffs(
     return PTThermalCoeffs(rheology, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL)
 end
 
+function JR2D.update_thermal_coeffs!(
+    pt_thermal::JustRelax.PTThermalCoeffs{T,<:CuArray}, rheology, phase_ratios, args, dt
+) where {T}
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ,
+        pt_thermal.dτ_ρ,
+        rheology,
+        phase_ratios.center,
+        args,
+        pt_thermal.max_lxyz,
+        pt_thermal.Vpdτ,
+        inv(dt),
+    )
+    return nothing
+end
+
+function JR2D.update_thermal_coeffs!(
+    pt_thermal::JustRelax.PTThermalCoeffs{T,<:CuArray}, rheology, args, dt
+) where {T}
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ,
+        pt_thermal.dτ_ρ,
+        rheology,
+        args,
+        pt_thermal.max_lxyz,
+        pt_thermal.Vpdτ,
+        inv(dt),
+    )
+    return nothing
+end
+
+function JR2D.update_thermal_coeffs!(
+    pt_thermal::JustRelax.PTThermalCoeffs{T,<:CuArray}, rheology, ::Nothing, args, dt
+) where {T}
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ,
+        pt_thermal.dτ_ρ,
+        rheology,
+        args,
+        pt_thermal.max_lxyz,
+        pt_thermal.Vpdτ,
+        inv(dt),
+    )
+    return nothing
+end
+
 # Boundary conditions
 function JR2D.flow_bcs!(
     ::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions
diff --git a/src/ext/CUDA/3D.jl b/src/ext/CUDA/3D.jl
index 7013245a..dfbdd21d 100644
--- a/src/ext/CUDA/3D.jl
+++ b/src/ext/CUDA/3D.jl
@@ -95,6 +95,55 @@ function JR3D.update_pt_thermal_arrays!(
     return nothing
 end
 
+function JR3D.update_thermal_coeffs!(
+    pt_thermal::JustRelax.PTThermalCoeffs{T,<:CuArray}, rheology, phase_ratios, args, dt
+) where {T}
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ,
+        pt_thermal.dτ_ρ,
+        rheology,
+        phase_ratios.center,
+        args,
+        pt_thermal.max_lxyz,
+        pt_thermal.Vpdτ,
+        inv(dt),
+    )
+    return nothing
+end
+
+function JR3D.update_thermal_coeffs!(
+    pt_thermal::JustRelax.PTThermalCoeffs{T,<:CuArray}, rheology, args, dt
+) where {T}
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ,
+        pt_thermal.dτ_ρ,
+        rheology,
+        args,
+        pt_thermal.max_lxyz,
+        pt_thermal.Vpdτ,
+        inv(dt),
+    )
+    return nothing
+end
+
+function JR3D.update_thermal_coeffs!(
+    pt_thermal::JustRelax.PTThermalCoeffs{T,<:CuArray}, rheology, ::Nothing, args, dt
+) where {T}
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ,
+        pt_thermal.dτ_ρ,
+        rheology,
+        args,
+        pt_thermal.max_lxyz,
+        pt_thermal.Vpdτ,
+        inv(dt),
+    )
+    return nothing
+end
+
 # Boundary conditions
 function JR3D.flow_bcs!(
     ::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions
diff --git a/src/thermal_diffusion/DiffusionPT_coefficients.jl b/src/thermal_diffusion/DiffusionPT_coefficients.jl
index acd809c0..2ee4f643 100644
--- a/src/thermal_diffusion/DiffusionPT_coefficients.jl
+++ b/src/thermal_diffusion/DiffusionPT_coefficients.jl
@@ -121,3 +121,50 @@ function _compute_pt_thermal_arrays!(
 
     return nothing
 end
+
+function update_thermal_coeffs!(
+    pt_thermal::JustRelax.PTThermalCoeffs, rheology, phase_ratios, args, dt
+)
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ,
+        pt_thermal.dτ_ρ,
+        rheology,
+        phase_ratios.center,
+        args,
+        pt_thermal.max_lxyz,
+        pt_thermal.Vpdτ,
+        inv(dt),
+    )
+    return nothing
+end
+
+function update_thermal_coeffs!(pt_thermal::JustRelax.PTThermalCoeffs, rheology, args, dt)
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ,
+        pt_thermal.dτ_ρ,
+        rheology,
+        args,
+        pt_thermal.max_lxyz,
+        pt_thermal.Vpdτ,
+        inv(dt),
+    )
+    return nothing
+end
+
+function update_thermal_coeffs!(
+    pt_thermal::JustRelax.PTThermalCoeffs, rheology, ::Nothing, args, dt
+)
+    ni = size(pt_thermal.dτ_ρ)
+    @parallel (@idx ni) compute_pt_thermal_arrays!(
+        pt_thermal.θr_dτ,
+        pt_thermal.dτ_ρ,
+        rheology,
+        args,
+        pt_thermal.max_lxyz,
+        pt_thermal.Vpdτ,
+        inv(dt),
+    )
+    return nothing
+end
diff --git a/src/thermal_diffusion/DiffusionPT_kernels.jl b/src/thermal_diffusion/DiffusionPT_kernels.jl
index 93b03fed..281f42b2 100644
--- a/src/thermal_diffusion/DiffusionPT_kernels.jl
+++ b/src/thermal_diffusion/DiffusionPT_kernels.jl
@@ -127,14 +127,15 @@ end
     d_za(A) = _d_za(A, I..., _dz)
 
     I1 = I .+ 1
-    T[I1...] +=
-        av(dτ_ρ) * (
-            (-(d_xa(qTx) + d_ya(qTy) + d_za(qTz))) -
-            av(ρCp) * (T[I1...] - Told[I1...]) * _dt
-        ) +
-        av(H) +
-        av(shear_heating)
-
+    T[I1...] =
+        (
+            av(dτ_ρ) * (
+                -(d_xa(qTx) + d_ya(qTy) + d_za(qTz)) +
+                Told[I1...] * av(ρCp) * _dt +
+                av(H) +
+                av(shear_heating)
+            ) + T[I1...]
+        ) / (one(_T) + av(dτ_ρ) * av(ρCp) * _dt)
     return nothing
 end
 
@@ -166,16 +167,18 @@ end
     T_ijk = T[I...]
     args_ijk = (; T=T_ijk, P=av(args.P))
     phase_ijk = getindex_phase(phase, i, j, k)
+    ρCp = compute_ρCp(rheology, phase_ijk, args_ijk)
 
     T[I...] =
-        T_ijk +
-        av(dτ_ρ) * (
-            (-(d_xa(qTx) + d_ya(qTy) + d_za(qTz))) -
-            compute_ρCp(rheology, phase_ijk, args_ijk) * (T_ijk - Told[I...]) * _dt +
-            av(H) +
-            av(shear_heating) +
-            adiabatic[i, j, k] * T_ijk
-        )
+        (
+            av(dτ_ρ) * (
+                -(d_xa(qTx) + d_ya(qTy) + d_za(qTz)) +
+                Told[I...] * ρCp * _dt +
+                av(H) +
+                av(shear_heating) +
+                adiabatic[i, j] * T_ijk
+            ) + T_ijk
+        ) / (one(_T) + av(dτ_ρ) * ρCp * _dt)
     return nothing
 end
 
@@ -389,13 +392,16 @@ end
     end
     #! format: on
 
-    T[i + 1, j + 1] +=
-        av(dτ_ρ) * (
-            (-(d_xa(qTx) + d_ya(qTy))) -
-            av(ρCp) * (T[i + 1, j + 1] - Told[i + 1, j + 1]) * _dt +
-            av(H) +
-            av(shear_heating)
-        )
+    T[i + 1, j + 1] =
+        (
+            av(dτ_ρ) * (
+                -(d_xa(qTx) + d_ya(qTy)) +
+                Told[i + 1, j + 1] * av(ρCp) * _dt +
+                av(H) +
+                av(shear_heating)
+            ) + T[i + 1, j + 1]
+        ) / (one(_T) + av(dτ_ρ) * av(ρCp) * _dt)
+
     return nothing
 end
 
@@ -437,13 +443,17 @@ end
             compute_ρCp(rheology, getindex_phase(phase, i1, j1), args_ij)
         ) * 0.25
 
-    T[i + 1, j + 1] +=
-        av(dτ_ρ) * (
-            (-(d_xa(qTx) + d_ya(qTy))) - ρCp * (T_ij - Told[i + 1, j + 1]) * _dt +
-            av(H) +
-            av(shear_heating) +
-            adiabatic[i, j] * T[i + 1, j + 1]
-        )
+    T[i + 1, j + 1] =
+        (
+            av(dτ_ρ) * (
+                -(d_xa(qTx) + d_ya(qTy)) +
+                Told[i + 1, j + 1] * ρCp * _dt +
+                av(H) +
+                av(shear_heating) +
+                adiabatic[i, j] * T[i + 1, j + 1]
+            ) + T[i + 1, j + 1]
+        ) / (one(_T) + av(dτ_ρ) * ρCp * _dt)
+
     return nothing
 end
 
diff --git a/src/thermal_diffusion/DiffusionPT_solver.jl b/src/thermal_diffusion/DiffusionPT_solver.jl
index dc4b76df..5f2bb886 100644
--- a/src/thermal_diffusion/DiffusionPT_solver.jl
+++ b/src/thermal_diffusion/DiffusionPT_solver.jl
@@ -147,6 +147,7 @@ function _heatdiffusion_PT!(
 
     while err > ϵ && iter < iterMax
         wtime0 += @elapsed begin
+            update_thermal_coeffs!(pt_thermal, rheology, phase, args, dt)
             @parallel flux_range(ni...) compute_flux!(
                 @qT(thermal)...,
                 @qT2(thermal)...,
diff --git a/test/test_diffusion2D.jl b/test/test_diffusion2D.jl
index d3891829..e9354188 100644
--- a/test/test_diffusion2D.jl
+++ b/test/test_diffusion2D.jl
@@ -74,7 +74,7 @@ function diffusion_2D(; nx=32, ny=32, lx=100e3, ly=100e3, ρ0=3.3e3, Cp0=1.2e3,
     )
     # fields needed to compute density on the fly
     P          = @zeros(ni...)
-    args       = (; P=P)
+    args       = (; P=P, T=@zeros(ni.+1...))
 
     ## Allocate arrays needed for every Thermal Diffusion
     thermal    = ThermalArrays(backend_JR, ni)
@@ -85,7 +85,7 @@ function diffusion_2D(; nx=32, ny=32, lx=100e3, ly=100e3, ρ0=3.3e3, Cp0=1.2e3,
     K          = @fill(K0, ni...)
     ρCp        = @. Cp * ρ
 
-    pt_thermal = PTThermalCoeffs(backend_JR, K, ρCp, dt, di, li)
+    pt_thermal = PTThermalCoeffs(backend_JR, K, ρCp, dt, di, li; CFL = 0.95 / √2.1)
     thermal_bc = TemperatureBoundaryConditions(;
         no_flux = (left = true, right = true, top = false, bot = false),
     )
diff --git a/test/test_diffusion2D_multiphase.jl b/test/test_diffusion2D_multiphase.jl
index 49ec80e7..cc39b725 100644
--- a/test/test_diffusion2D_multiphase.jl
+++ b/test/test_diffusion2D_multiphase.jl
@@ -165,7 +165,7 @@ function diffusion_2D(; nx=32, ny=32, lx=100e3, ly=100e3, Cp0=1.2e3, K0=3.0)
     # PT coefficients for thermal diffusion
     args       = (; P=P, T=thermal.Tc)
     pt_thermal = PTThermalCoeffs(
-        backend_JR, rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=0.65 / √2
+        backend_JR, rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=0.95 / √2
     )
 
     # Time loop
diff --git a/test/test_diffusion3D.jl b/test/test_diffusion3D.jl
index 41ea255f..ab86c624 100644
--- a/test/test_diffusion3D.jl
+++ b/test/test_diffusion3D.jl
@@ -89,7 +89,7 @@ function diffusion_3D(;
 
     # fields needed to compute density on the fly
     P          = @zeros(ni...)
-    args       = (; P=P)
+    args       = (; P=P, T=@zeros(ni.+1...))
 
     ## Allocate arrays needed for every Thermal Diffusion
     # general thermal arrays
@@ -102,7 +102,7 @@ function diffusion_3D(;
     ρCp        = @. Cp * ρ
 
     # Boundary conditions
-    pt_thermal = PTThermalCoeffs(backend, K, ρCp, dt, di, li; CFL = 0.75 / √3.1)
+    pt_thermal = PTThermalCoeffs(backend, K, ρCp, dt, di, li; CFL = 0.95 / √3.1)
     thermal_bc = TemperatureBoundaryConditions(;
         no_flux     = (left = true , right = true , top = false, bot = false, front = true , back = true),
     )
diff --git a/test/test_diffusion3D_multiphase.jl b/test/test_diffusion3D_multiphase.jl
index 255486e2..13bfcd8d 100644
--- a/test/test_diffusion3D_multiphase.jl
+++ b/test/test_diffusion3D_multiphase.jl
@@ -177,7 +177,7 @@ function diffusion_3D(;
 
     # PT coefficients for thermal diffusion
     args       = (; P=P, T=thermal.Tc)
-    pt_thermal = PTThermalCoeffs(backend_JR, K, ρCp, dt, di, li; CFL = 0.75 / √3.1)
+    pt_thermal = PTThermalCoeffs(backend_JR, K, ρCp, dt, di, li; CFL = 0.95 / √3.1)
 
     t  = 0.0
     it = 0