From ac1fb86db3c3e30aeae4a7f2177b8ccb8e490507 Mon Sep 17 00:00:00 2001 From: Baudouin Raoult Date: Thu, 9 Jan 2025 16:58:38 +0000 Subject: [PATCH] update --- src/anemoi/inference/memory.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/anemoi/inference/memory.py b/src/anemoi/inference/memory.py index b3482cb..e3884c1 100644 --- a/src/anemoi/inference/memory.py +++ b/src/anemoi/inference/memory.py @@ -78,10 +78,9 @@ def __call__(self, title): bytes_to_human(torch.cuda.memory_reserved(0)), ) - free, total = torch.cuda.mem_get_info("cuda:0") - mem_used_MB = (total - free) / 1024**2 - LOG.info(f"Pytorch mem_get_info says 'cuda:0' is using {mem_used_MB}MB / {total / 1024 ** 2}MB of memory.") - self.last_total = total self.last_title = title self.known = tensors + + free, total = torch.cuda.mem_get_info("cuda:0") + LOG.info(f"Pytorch mem_get_info says 'cuda:0' is using {bytes_to_human(total - free)} of memory.")