From d554d437a7e00fa9272b07329a72b6ff3af92c0e Mon Sep 17 00:00:00 2001
From: Zhengxiao Du <zx-du20@mails.tsinghua.edu.cn>
Date: Sun, 21 Nov 2021 16:20:54 +0800
Subject: [PATCH] Fix report memory in utils.py
---
SwissArmyTransformer/training/utils.py | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/SwissArmyTransformer/training/utils.py b/SwissArmyTransformer/training/utils.py
index 066efac..7919762 100755
--- a/SwissArmyTransformer/training/utils.py
+++ b/SwissArmyTransformer/training/utils.py
@@ -126,8 +126,8 @@ def report_memory(name):
torch.cuda.memory_allocated() / mega_bytes)
string += ' | max allocated: {}'.format(
torch.cuda.max_memory_allocated() / mega_bytes)
- string += ' | cached: {}'.format(torch.cuda.memory_cached() / mega_bytes)
+ string += ' | cached: {}'.format(torch.cuda.memory_reserved() / mega_bytes)
string += ' | max cached: {}'.format(
- torch.cuda.memory_reserved() / mega_bytes)
+ torch.cuda.max_memory_reserved() / mega_bytes)
print_rank_0(string)
--
GitLab