change to reset_peak_memory_stats

Author: Felipe Mello

src/forge/observability/perf_tracker.py

@@ -186,7 +186,7 @@ def _start_memory_tracking(self) -> None:
 
         if should_track:
             _set_memory_active(True)
-            torch.cuda.reset_max_memory_allocated()
+            torch.cuda.reset_peak_memory_stats()
             self._start_mem = torch.cuda.memory_allocated()
             self._memory_started = True
 
@@ -202,7 +202,7 @@ def _stop_memory_tracking(self) -> None:
         )
         record_metric(f"{self.prefix}/memory_peak_max_gb", peak_mem, Reduce.MAX)
         _set_memory_active(False)
-        torch.cuda.reset_max_memory_allocated()
+        torch.cuda.reset_peak_memory_stats()
         self._memory_started = False
 
     def _record_timing_metrics(