deepspeed/runtime/utils.py: reset_peak_memory_stats when empty cache (#2803)

Co-authored-by: N Olatunji Ruwase <olruwase@microsoft.com>

deepspeed/runtime/utils.py: reset_peak_memory_stats when empty cache (#2803)
Co-authored-by: N Olatunji Ruwase <olruwase@microsoft.com>
6eca037c · 郭叶军 · GitHub · 4d27225f · 6eca037c
隐藏空白更改
内联并排

Showing with 1 addition and 0 deletion

deepspeed/runtime/utils.py deepspeed/runtime/utils.py +1 -0

未找到文件。
--- 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):