Move check mem usage inside trace call (#2715)

Summary:
Pull Request resolved: #2715

As title. Move the mem usage collection inside the trace call.

Reviewed By: gag1jain, satgera

Differential Revision: D68914259

fbshipit-source-id: bab5153f5b452497c468244eab6687a5e43d1460

Author: atuljangra

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()