Use ClimaComms.@time instead of @time

For GPU compatibility
CliMA · Apr 30, 2024 · f5897b9 · f5897b9
1 parent a643f22
commit f5897b9
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/experiments/standalone/Bucket/global_bucket_function.jl b/experiments/standalone/Bucket/global_bucket_function.jl
@@ -161,7 +161,7 @@ updatefunc = ClimaLand.make_update_drivers(bucket_atmos, bucket_rad)
 driver_cb = ClimaLand.DriverUpdateCallback(updateat, updatefunc)
 cb = SciMLBase.CallbackSet(driver_cb, saving_cb)
 
-ClimaComms.@time sol =
+ClimaComms.@time ClimaComms.device() sol =
     SciMLBase.solve(prob, ode_algo; dt = Δt, saveat = saveat, callback = cb);
 
 # Interpolate to grid

diff --git a/experiments/standalone/Bucket/global_bucket_staticmap.jl b/experiments/standalone/Bucket/global_bucket_staticmap.jl
@@ -176,7 +176,7 @@ updatefunc = ClimaLand.make_update_drivers(bucket_atmos, bucket_rad)
 driver_cb = ClimaLand.DriverUpdateCallback(updateat, updatefunc)
 cb = SciMLBase.CallbackSet(driver_cb, saving_cb)
 
-@time sol =
+ClimaComms.@time ClimaComms.device() sol =
     SciMLBase.solve(prob, ode_algo; dt = Δt, saveat = saveat, callback = cb);
 
 # Interpolate to grid

diff --git a/experiments/standalone/Bucket/global_bucket_temporalmap.jl b/experiments/standalone/Bucket/global_bucket_temporalmap.jl
@@ -194,7 +194,7 @@ prob, cb, saveat, saved_values = setup_prob(t0, tf, Δt);
 timestepper = CTS.RK4()
 ode_algo = CTS.ExplicitAlgorithm(timestepper)
 
-@time sol =
+ClimaComms.@time ClimaComms.device() sol =
     SciMLBase.solve(prob, ode_algo; dt = Δt, saveat = saveat, callback = cb)
 if PROFILING
     # Now that we compiled, solve again but collect profiling information