fix node manipulation in partition class (#48016)
Summary:
This PR fixes the add_node and remove_node in partition class and also add a unit test for node manipulation in partition
Pull Request resolved: https://github.com/pytorch/pytorch/pull/48016
Reviewed By: gcatron
Differential Revision: D24996368
Pulled By: scottxu0730
fbshipit-source-id: 0ddffd5ed3f95e5285fffcaee8c4b671929b4df3
