[fx-acc] Fix wrong device assignment in find_single_partition (#60056)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/60056
Previously we put the whole graph as a single partition onto a device with maximum memory if possible, but the code assumed that the first logical device always has the maximum memory.
This diff fixes this issue and updates the unittest to reflect such a corner case.
Test Plan:
```
buck test mode/opt //caffe2/test:test_fx_experimental -- --exact 'caffe2/test:test_fx_experimental - test_find_single_partition (test_fx_experimental.TestFXExperimental)'
Started reporting to test run: https://www.internalfb.com/intern/testinfra/testrun/6473924507772744
✓ ListingSuccess: caffe2/test:test_fx_experimental - main (1.357)
✓ Pass: caffe2/test:test_fx_experimental - test_find_single_partition (test_fx_experimental.TestFXExperimental) (1.206)
Summary
Pass: 1
ListingSuccess: 1
```
Reviewed By: gcatron
Differential Revision: D29118715
fbshipit-source-id: cac6a1f0d2f47717446dcc80093bbcf362663859
