From 93324c11fbf40043056ce120abacf9798611bb3a Mon Sep 17 00:00:00 2001
From: Charles Kawczynski <kawczynski.charles@gmail.com>
Date: Mon, 17 Jul 2023 12:37:00 -0700
Subject: [PATCH] Use ClimaTimeSteppers in CUDA shallow water example

---
 examples/sphere/shallow_water_cuda.jl | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/examples/sphere/shallow_water_cuda.jl b/examples/sphere/shallow_water_cuda.jl
index 71fd90637b..20d2dcd8fd 100644
--- a/examples/sphere/shallow_water_cuda.jl
+++ b/examples/sphere/shallow_water_cuda.jl
@@ -3,6 +3,8 @@ using ClimaComms
 using DocStringExtensions
 using LinearAlgebra
 using ClimaTimeSteppers, DiffEqBase
+import OrdinaryDiffEq as ODE
+import ClimaTimeSteppers as CTS
 using DiffEqCallbacks
 using NVTX
 
@@ -564,12 +566,15 @@ function shallow_water_driver_cuda(ARGS, ::Type{FT}) where {FT}
         end
         return nothing
     end
-
-    prob =
-        ODEProblem(ClimaODEFunction(; T_exp! = rhs!), Y, (0.0, T), parameters)
-    integrator = DiffEqBase.init(
+    prob = ODE.ODEProblem(
+        CTS.ClimaODEFunction(; T_exp! = rhs!),
+        Y,
+        (FT(0), T),
+        parameters,
+    )
+    integrator = ODE.init(
         prob,
-        ExplicitAlgorithm(SSP33ShuOsher()),
+        CTS.ExplicitAlgorithm(CTS.SSP33ShuOsher()),
         dt = dt,
         saveat = [],
         progress = true,