fix: always show total last

brycedrennan · Jan 2, 2024 · e2b3e82 · e2b3e82
1 parent b97a8a2
commit e2b3e82
Showing 1 changed file with 3 additions and 0 deletions.
diff --git 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)