Fix profiler memory logging (#2735)

mosaicml · Nov 21, 2023 · a193115 · a193115
1 parent 92d983a
commit a193115
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/composer/profiler/torch_profiler.py b/composer/profiler/torch_profiler.py
@@ -256,9 +256,11 @@ def close(self, state: State, logger: Logger) -> None:
         del state, logger  # unused
         if self.profiler is not None:
             log.info(self.profiler.key_averages().table(sort_by='cpu_time_total', row_limit=20))
-            log.info(self.profiler.key_averages().table(sort_by='self_cpu_memory_usage', row_limit=20))
+            if self.profile_memory:
+                log.info(self.profiler.key_averages().table(sort_by='self_cpu_memory_usage', row_limit=20))
             if torch.profiler.ProfilerActivity.CUDA in self.profiler.activities:
                 log.info(self.profiler.key_averages().table(sort_by='cuda_time_total', row_limit=20))
-                log.info(self.profiler.key_averages().table(sort_by='self_cuda_memory_usage', row_limit=20))
+                if self.profile_memory:
+                    log.info(self.profiler.key_averages().table(sort_by='self_cuda_memory_usage', row_limit=20))
             self.profiler.__exit__(None, None, None)
             self.profiler = None