[PyTorch] Rework stat collection in CUDACachingAllocator (#71669)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/71669
This was relatively inefficient. Rather than looping for each type of stat we want to update, we now do one loop covering all the stats.
ghstack-source-id: 148013645
Reviewed By: ngimel
Differential Revision: D33725458
fbshipit-source-id: 39ef5d65a73d4ef67f259de8c02c7df29487d990
(cherry picked from commit 7ca46689b72ba7611517447a292445571bd02dd7)
