Fix mis-calculated memory compression ratio (#150695)
Summary:
https://github.com/pytorch/pytorch/pull/149817 introduced an extra warmup run to compute AOTI memory compression ratio, but since weights are only loaded once in the AOTI run, the peak memory seen in the extra warmup won't include the weight, which causes an aritifically high memory compression ratio. This PR removes that extra warmup run, and calls reset_peak_memory_stats in the proper place instead.
X-link: https://github.com/pytorch/pytorch/pull/150695
Approved by: https://github.com/yushangdi
Reviewed By: atalman
Differential Revision: D72570207
fbshipit-source-id: 421bde2de6ebde4ca795871c5920a66a6b77073f
