diff --git a/experiments/AMIP/coupler_driver.jl b/experiments/AMIP/coupler_driver.jl index 44e02809c3..87412f87e1 100644 --- 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