diff --git a/deepspeed/runtime/utils.py b/deepspeed/runtime/utils.py index bbdbee1924528d96e4eda69b1d64284c8d20302c..95a0b9621153d3c5db275ba7e14c3b7009a83968 100755 --- a/deepspeed/runtime/utils.py +++ b/deepspeed/runtime/utils.py @@ -765,6 +765,7 @@ def get_ma_status(): def empty_cache(): get_accelerator().empty_cache() + get_accelerator().reset_peak_memory_stats() def see_memory_usage(message, force=False):