Develop fix max_mem_reserved for benchmark (#829)

* fix profile * add max_mem_reserved for benchmark * fix
PaddlePaddle · Dec 5, 2023 · 21ee373 · 21ee373
1 parent c7f961f
commit 21ee373
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/ppgan/engine/trainer.py b/ppgan/engine/trainer.py
@@ -355,10 +355,11 @@ def print_log(self):
             eta_str = str(datetime.timedelta(seconds=int(eta)))
             message += f'eta: {eta_str}'
 
-        max_mem_reserved_str = f" max_mem_reserved: {paddle.device.cuda.max_memory_reserved()} B"
-        max_mem_allocated_str = f" max_mem_allocated: {paddle.device.cuda.max_memory_allocated()} B"
-        message += max_mem_reserved_str
-        message += max_mem_allocated_str
+        if paddle.device.is_compiled_with_cuda():
+            max_mem_reserved_str = f" max_mem_reserved: {paddle.device.cuda.max_memory_reserved() // (1024 ** 2)} MB"
+            max_mem_allocated_str = f" max_mem_allocated: {paddle.device.cuda.max_memory_allocated() // (1024 ** 2)} MB"
+            message += max_mem_reserved_str
+            message += max_mem_allocated_str
         # print the message
         self.logger.info(message)