From e2b3e82cae50c6f902f1a91646c8b6fe75730272 Mon Sep 17 00:00:00 2001 From: Bryce Date: Mon, 1 Jan 2024 18:58:41 -0800 Subject: [PATCH] fix: always show total last --- imaginairy/utils/log_utils.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/imaginairy/utils/log_utils.py b/imaginairy/utils/log_utils.py index f294ca0d..0b55cb39 100644 --- a/imaginairy/utils/log_utils.py +++ b/imaginairy/utils/log_utils.py @@ -179,6 +179,9 @@ def get_performance_stats(self) -> dict[str, dict[str, float]]: self.summary_context.stop() self.timing_contexts["total"] = self.summary_context + # move total to the end + self.timing_contexts["total"] = self.timing_contexts.pop("total") + if torch.cuda.is_available(): self.summary_context.memory_peak = max( max(context.memory_peak, context.memory_start, context.memory_end)