cleanups to memory_dag (#21430)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/21430
ghimport-source-id: 2dc5a0df8512e796c12d65d3ecc5981638122ce6
Reviewed By: jamesr66a
Differential Revision: D15718561
Pulled By: suo
fbshipit-source-id: 1ef31c08c8a757b632451eb07a47a8227e76c67f
