From b99146fd863b01794c602c6fd3103ce8dc609928 Mon Sep 17 00:00:00 2001 From: Atul Jangra Date: Thu, 30 Jan 2025 13:27:46 -0800 Subject: [PATCH] Move check mem usage inside trace call (#2715) Summary: As title. Move the mem usage collection inside the trace call. Reviewed By: gag1jain, satgera Differential Revision: D68914259 --- torchrec/metrics/metric_module.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchrec/metrics/metric_module.py b/torchrec/metrics/metric_module.py index 56356ff8b..e4417e2c0 100644 --- a/torchrec/metrics/metric_module.py +++ b/torchrec/metrics/metric_module.py @@ -353,8 +353,8 @@ def compute(self) -> Dict[str, MetricValue]: right before logging the metrics results to the data sink. """ self.compute_count += 1 - self.check_memory_usage(self.compute_count) with record_function("## RecMetricModule:compute ##"): + self.check_memory_usage(self.compute_count) ret: Dict[str, MetricValue] = {} if self.rec_metrics: self._adjust_compute_interval()