Move check mem usage inside trace call

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

Reviewed By: gag1jain

Differential Revision: D68914259

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