From 6e528b4a8263fdc2bec4a00df9625ae0bd79bf30 Mon Sep 17 00:00:00 2001
From: Julia Sloan <jsloan@caltech.edu>
Date: Mon, 10 Jun 2024 21:18:18 -0700
Subject: [PATCH] ClimaComms @time fix

---
 experiments/standalone/Bucket/global_bucket_function.jl  | 9 +++++++--
 experiments/standalone/Bucket/global_bucket_staticmap.jl | 9 +++++++--
 .../standalone/Bucket/global_bucket_temporalmap.jl       | 9 +++++++--
 experiments/standalone/Soil/richards_runoff.jl           | 2 +-
 4 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/experiments/standalone/Bucket/global_bucket_function.jl b/experiments/standalone/Bucket/global_bucket_function.jl
index b098c23b5e..fae2f9111b 100644
--- a/experiments/standalone/Bucket/global_bucket_function.jl
+++ b/experiments/standalone/Bucket/global_bucket_function.jl
@@ -164,8 +164,13 @@ updatefunc = ClimaLand.make_update_drivers(bucket_atmos, bucket_rad)
 driver_cb = ClimaLand.DriverUpdateCallback(updateat, updatefunc)
 cb = SciMLBase.CallbackSet(driver_cb, saving_cb)
 
-ClimaComms.@time ClimaComms.device() sol =
-    SciMLBase.solve(prob, ode_algo; dt = Δt, saveat = saveat, callback = cb);
+sol = ClimaComms.@time ClimaComms.device() SciMLBase.solve(
+    prob,
+    ode_algo;
+    dt = Δt,
+    saveat = saveat,
+    callback = cb,
+);
 
 # Interpolate to grid
 space = axes(coords.surface)
diff --git a/experiments/standalone/Bucket/global_bucket_staticmap.jl b/experiments/standalone/Bucket/global_bucket_staticmap.jl
index 6921713af1..cbf07ace4c 100644
--- a/experiments/standalone/Bucket/global_bucket_staticmap.jl
+++ b/experiments/standalone/Bucket/global_bucket_staticmap.jl
@@ -247,8 +247,13 @@ updatefunc = ClimaLand.make_update_drivers(bucket_atmos, bucket_rad)
 driver_cb = ClimaLand.DriverUpdateCallback(updateat, updatefunc)
 cb = SciMLBase.CallbackSet(driver_cb, saving_cb)
 
-ClimaComms.@time ClimaComms.device() sol =
-    SciMLBase.solve(prob, ode_algo; dt = Δt, saveat = saveat, callback = cb);
+sol = ClimaComms.@time ClimaComms.device() SciMLBase.solve(
+    prob,
+    ode_algo;
+    dt = Δt,
+    saveat = saveat,
+    callback = cb,
+);
 
 # Interpolate to grid
 space = axes(coords.surface)
diff --git a/experiments/standalone/Bucket/global_bucket_temporalmap.jl b/experiments/standalone/Bucket/global_bucket_temporalmap.jl
index bfb72ebcf6..05cee96a01 100644
--- a/experiments/standalone/Bucket/global_bucket_temporalmap.jl
+++ b/experiments/standalone/Bucket/global_bucket_temporalmap.jl
@@ -192,8 +192,13 @@ prob, cb, saveat, saved_values = setup_prob(t0, tf, Δt);
 timestepper = CTS.RK4()
 ode_algo = CTS.ExplicitAlgorithm(timestepper)
 
-ClimaComms.@time ClimaComms.device() sol =
-    SciMLBase.solve(prob, ode_algo; dt = Δt, saveat = saveat, callback = cb)
+sol = ClimaComms.@time ClimaComms.device() SciMLBase.solve(
+    prob,
+    ode_algo;
+    dt = Δt,
+    saveat = saveat,
+    callback = cb,
+)
 if PROFILING
     # Now that we compiled, solve again but collect profiling information
 
diff --git a/experiments/standalone/Soil/richards_runoff.jl b/experiments/standalone/Soil/richards_runoff.jl
index 53750d8275..53430d655d 100644
--- a/experiments/standalone/Soil/richards_runoff.jl
+++ b/experiments/standalone/Soil/richards_runoff.jl
@@ -256,7 +256,7 @@ updateat = Array(t0:dt:tf)
 updatefunc = ClimaLand.make_update_drivers(atmos, nothing)
 driver_cb = ClimaLand.DriverUpdateCallback(updateat, updatefunc)
 cb = SciMLBase.CallbackSet(driver_cb, saving_cb)
-@time sol = SciMLBase.solve(prob, ode_algo; dt = dt, saveat = dt, callback = cb)
+sol = @time SciMLBase.solve(prob, ode_algo; dt = dt, saveat = dt, callback = cb)
 
 # Make plots on CPU
 if context.device isa ClimaComms.CPUSingleThreaded