add get_device_to_partitions_mapping (#47361)
Summary:
add get_device_to_partitions_mapping function in the Partitioner class to make size_based_partition more modular and organized. This function will also be used in the future cost_aware_partition
Pull Request resolved: https://github.com/pytorch/pytorch/pull/47361
Reviewed By: gcatron
Differential Revision: D24760911
Pulled By: scottxu0730
fbshipit-source-id: 8cdda51b9a1145f9d13ebabbb98b4d9df5ebb6cd
