model_dump: Fix non-counting and double-counting bugs in tensor memory (#60702)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/60702
- Instead of traversing and counting all tensor memory, collect a map
from storage key to storage info while traversing. Add up sizes at
the end to avoid double counting.
- Count tensor memory from constants as well.
Test Plan: Ran webdriver test.
Reviewed By: dhruvbird
Differential Revision: D29380396
Pulled By: dreiss
fbshipit-source-id: 6d0fd66f677fe23c851aa218387aa4dc59502b1e
