diff --git a/src/JustRelax_CPU.jl b/src/JustRelax_CPU.jl
index a8dfbf18..5add960f 100644
--- a/src/JustRelax_CPU.jl
+++ b/src/JustRelax_CPU.jl
@@ -18,8 +18,9 @@ import JustRelax:
     AbstractFlowBoundaryConditions,
     DisplacementBoundaryConditions,
     VelocityBoundaryConditions,
-    apply_dirichlet, apply_dirichlet!
-    
+    apply_dirichlet,
+    apply_dirichlet!
+
 import JustPIC._2D: numphases, nphases
 
 __init__() = @init_parallel_stencil(Threads, Float64, 2)
@@ -49,8 +50,9 @@ import JustRelax:
     TemperatureBoundaryConditions,
     AbstractFlowBoundaryConditions,
     DisplacementBoundaryConditions,
-    VelocityBoundaryConditions, 
-    apply_dirichlet, apply_dirichlet!
+    VelocityBoundaryConditions,
+    apply_dirichlet,
+    apply_dirichlet!
 
 import JustPIC._3D: numphases, nphases
 
diff --git a/src/boundaryconditions/Dirichlet.jl b/src/boundaryconditions/Dirichlet.jl
index 361aaa67..0c8d259f 100644
--- a/src/boundaryconditions/Dirichlet.jl
+++ b/src/boundaryconditions/Dirichlet.jl
@@ -71,49 +71,49 @@ function Base.getindex(
 end
 
 @inline function apply_dirichlet!(A::AbstractArray, bc::AbstractDirichletBoundaryCondition)
-    apply_mask!(A, bc.value, bc.mask)
+    return apply_mask!(A, bc.value, bc.mask)
 end
 
 @inline function apply_dirichlet!(
     ::AbstractArray, ::AbstractDirichletBoundaryCondition{Nothing,Nothing}
-) 
-    nothing
+)
+    return nothing
 end
 
 @inline function apply_dirichlet!(
     A::AbstractArray, bc::AbstractDirichletBoundaryCondition, inds::Vararg{Int,N}
-) where {N} 
-    apply_mask!(A, bc.value, bc.mask, inds...)
+) where {N}
+    return apply_mask!(A, bc.value, bc.mask, inds...)
 end
 
 @inline function apply_dirichlet!(
     ::AbstractArray, ::AbstractDirichletBoundaryCondition{Nothing,Nothing}, ::Vararg{Int,N}
-) where {N} 
-    nothing
+) where {N}
+    return nothing
 end
 
 @inline function apply_dirichlet(A::AbstractArray, bc::AbstractDirichletBoundaryCondition)
-    apply_mask(A, bc.value, bc.mask)
+    return apply_mask(A, bc.value, bc.mask)
 end
 
 @inline function apply_dirichlet(
     A::AbstractArray, ::AbstractDirichletBoundaryCondition{Nothing,Nothing}
-) 
- A
+)
+    return A
 end
 
 @inline function apply_dirichlet(
     A::AbstractArray, bc::AbstractDirichletBoundaryCondition, inds::Vararg{Int,N}
-) where {N} 
-    apply_mask(A, bc.value, bc.mask, inds...)
+) where {N}
+    return apply_mask(A, bc.value, bc.mask, inds...)
 end
 
 @inline function apply_dirichlet(
     A::AbstractArray,
     ::AbstractDirichletBoundaryCondition{Nothing,Nothing},
     inds::Vararg{Int,N},
-) where {N} 
-    A[inds...]
+) where {N}
+    return A[inds...]
 end
 
 @inline Dirichlet(x::NamedTuple) = Dirichlet(; x...)
diff --git a/src/mask/constructors.jl b/src/mask/constructors.jl
index 40a1080a..27269b2b 100644
--- a/src/mask/constructors.jl
+++ b/src/mask/constructors.jl
@@ -1,4 +1,4 @@
-function Mask(nx, ny, inds::Vararg{UnitRange, 2})
+function Mask(nx, ny, inds::Vararg{UnitRange,2})
     mask = @zeros(nx, ny)
     @views mask[inds...] .= 1
     return JustRelax.Mask(mask)
diff --git a/src/thermal_diffusion/DiffusionPT_kernels.jl b/src/thermal_diffusion/DiffusionPT_kernels.jl
index c47c68c3..dd4b7b0d 100644
--- a/src/thermal_diffusion/DiffusionPT_kernels.jl
+++ b/src/thermal_diffusion/DiffusionPT_kernels.jl
@@ -137,8 +137,8 @@ end
                 av(shear_heating)
             ) + T[I1...]
         ) / (one(_T) + av(dτ_ρ) * av(ρCp) * _dt)
-    apply_dirichlet!(T , dirichlet, I1...)
-    
+    apply_dirichlet!(T, dirichlet, I1...)
+
     return nothing
 end
 
@@ -183,7 +183,7 @@ end
                 adiabatic[i, j, k] * T_ijk
             ) + T_ijk
         ) / (one(_T) + av(dτ_ρ) * ρCp * _dt)
-    apply_dirichlet!(T , dirichlet, I1...)
+    apply_dirichlet!(T, dirichlet, I1...)
 
     return nothing
 end
@@ -388,7 +388,18 @@ end
 end
 
 @parallel_indices (i, j) function update_T!(
-    T::AbstractArray{_T,2}, Told, qTx, qTy, H, shear_heating, ρCp, dτ_ρ, dirichlet, _dt, _dx, _dy
+    T::AbstractArray{_T,2},
+    Told,
+    qTx,
+    qTy,
+    H,
+    shear_heating,
+    ρCp,
+    dτ_ρ,
+    dirichlet,
+    _dt,
+    _dx,
+    _dy,
 ) where {_T}
     nx, ny = size(ρCp)
 
@@ -405,7 +416,7 @@ end
     end
     #! format: on
 
-    I1  = i + 1, j + 1
+    I1 = i + 1, j + 1
     T[I1...] =
         (
             av(dτ_ρ) * (
@@ -416,8 +427,8 @@ end
             ) + T[I1...]
         ) / (one(_T) + av(dτ_ρ) * av(ρCp) * _dt)
 
-    apply_dirichlet!(T , dirichlet, I1...)
-    
+    apply_dirichlet!(T, dirichlet, I1...)
+
     return nothing
 end
 
@@ -460,7 +471,7 @@ end
             compute_ρCp(rheology, getindex_phase(phase, i1, j1), args_ij)
         ) * 0.25
 
-    I1  = i + 1, j + 1
+    I1 = i + 1, j + 1
     T[I1...] =
         (
             av(dτ_ρ) * (
@@ -471,14 +482,25 @@ end
                 adiabatic[i, j] * T[I1...]
             ) + T[I1...]
         ) / (one(_T) + av(dτ_ρ) * ρCp * _dt)
-    
-    apply_dirichlet!(T , dirichlet, I1...)
+
+    apply_dirichlet!(T, dirichlet, I1...)
 
     return nothing
 end
 
 @parallel_indices (i, j) function check_res!(
-    ResT::AbstractArray{_T,2}, T, Told, qTx2, qTy2, H, shear_heating, ρCp, dirichlet, _dt, _dx, _dy
+    ResT::AbstractArray{_T,2},
+    T,
+    Told,
+    qTx2,
+    qTy2,
+    H,
+    shear_heating,
+    ρCp,
+    dirichlet,
+    _dt,
+    _dx,
+    _dy,
 ) where {_T}
     nx, ny = size(ρCp)
 
@@ -497,14 +519,13 @@ end
 
     I1 = i + 1, j + 1
     ResT[i, j] = if iszero(dirichlet.mask[I1...])
-        -av(ρCp) * (T[I1...] - Told[I1...]) * _dt -
-        (d_xa(qTx2) + d_ya(qTy2)) +
+        -av(ρCp) * (T[I1...] - Told[I1...]) * _dt - (d_xa(qTx2) + d_ya(qTy2)) +
         av(H) +
         av(shear_heating)
     else
         zero(_T)
     end
-    
+
     return nothing
 end
 
@@ -553,7 +574,7 @@ end
         av(H) +
         av(shear_heating) +
         adiabatic[i, j] * T[I1...]
-    else 
+    else
         zero(_T)
     end
 
diff --git a/src/thermal_diffusion/DiffusionPT_solver.jl b/src/thermal_diffusion/DiffusionPT_solver.jl
index 9ecebecf..14c87245 100644
--- a/src/thermal_diffusion/DiffusionPT_solver.jl
+++ b/src/thermal_diffusion/DiffusionPT_solver.jl
@@ -53,7 +53,17 @@ function _heatdiffusion_PT!(
             @parallel flux_range(ni...) compute_flux!(
                 @qT(thermal)..., @qT2(thermal)..., thermal.T, K, pt_thermal.θr_dτ, _di...
             )
-            update_T(nothing, b_width, thermal, ρCp, pt_thermal, thermal_bc.dirichlet, _dt, _di, ni)
+            update_T(
+                nothing,
+                b_width,
+                thermal,
+                ρCp,
+                pt_thermal,
+                thermal_bc.dirichlet,
+                _dt,
+                _di,
+                ni,
+            )
             thermal_bcs!(thermal, thermal_bc)
             update_halo!(thermal.T)
         end
@@ -160,7 +170,17 @@ function _heatdiffusion_PT!(
                 args,
             )
             update_T(
-                nothing, b_width, thermal, rheology, phases, pt_thermal, thermal_bc.dirichlet, _dt, _di, ni, args
+                nothing,
+                b_width,
+                thermal,
+                rheology,
+                phases,
+                pt_thermal,
+                thermal_bc.dirichlet,
+                _dt,
+                _di,
+                ni,
+                args,
             )
             thermal_bcs!(thermal, thermal_bc)
             update_halo!(thermal.T)