fix: always show total last

This commit is contained in:
Bryce 2024-01-01 18:58:41 -08:00 committed by Bryce Drennan
parent f50a1f5b0c
commit 4fcfc363af

View File

@ -179,6 +179,9 @@ class ImageLoggingContext:
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)