fix show_mem_usage take

CliMA · Mar 22, 2024 · 581b0f9 · 581b0f9
1 parent 34c64c0
commit 581b0f9
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/experiments/AMIP/coupler_driver.jl b/experiments/AMIP/coupler_driver.jl
@@ -47,10 +47,14 @@ directly compared.
 function show_memory_usage(comms_ctx, objects, io::Union{Base.TTY, IOContext, IOBuffer} = stdout)
     if comms_ctx.device isa ClimaComms.CUDADevice
         # If `io` is `stdout`, print the memory status, otherwise store in buffer to return
-        CUDA.memory_status(io)
-        # Get the memory usage in GB by parsing the output of `CUDA.memory_status`
-        allocs_GB = split(split(take!(io), '(')[2], '/')[1]
-        return io != stdout ? allocs_GB : nothing
+        if io == stdout
+            CUDA.memory_status(io)
+            return nothing
+        else
+            # Get the memory usage in GB by parsing the output of `CUDA.memory_status`
+            allocs_GB = split(split(take!(io), '(')[2], '/')[1]
+            return allocs_GB
+        end
     elseif comms_ctx.device isa ClimaComms.AbstractCPUDevice
         if ClimaComms.iamroot(comms_ctx)
             cumul_size = 0