Fix split_module bug (#95493)
Summary: Title, the mapping currently has lots of unused keys due to the condition or always return True, but it will not affect the correctness.
Test Plan: N/A
Differential Revision: D43579510
Pull Request resolved: https://github.com/pytorch/pytorch/pull/95493
Approved by: https://github.com/Skylion007
