pytorch

a86b3438 - add support for different memory sizes on size_based_partition (#46919)

add support for different memory sizes on size_based_partition (#46919) Summary: WIP: add support for different memory sizes on size_based_partition, so the size_based_partition could support different logical devices with different memory sizes. Compared to the original size_based_partition, the new one also supports partition to logical device mapping. Multiple partitions can be mapped into one device if the memory size is allowed. A test unit test_different_size_partition is also added. Pull Request resolved: https://github.com/pytorch/pytorch/pull/46919 Reviewed By: gcatron, VitalyFedyunin Differential Revision: D24603511 Pulled By: scottxu0730 fbshipit-source-id: 1ba37338ae054ad846b425fbb7e631d3b6c500b6