mirror of
https://github.com/zebrajr/pytorch.git
synced 2026-01-15 12:15:51 +00:00
28ee8be5bf
Summary: Placement is leaked to too many classes! In this diff, we consolidate all placement lookup into one place: Graph::ApplyDevicePlacement. After applying placement, the in-memory graph, tensorMeta, weightMeta would already have the re-mapped device. The subsequence weight loading, sample input loading, target device inference would look up the re-mapped device from graph's tensorMeta. graph's tensorMeta becomes the only ground truth! Test Plan: Need to add some tests before landing. This is a big change. Rollback Plan: Differential Revision: D78841818 Pull Request resolved: https://github.com/pytorch/pytorch/pull/158996 Approved by: https://github.com/henryoier