torchrec icon indicating copy to clipboard operation
torchrec copied to clipboard

Implement unlift_and_swap function

Open angelayi opened this issue 1 year ago • 2 comments

Summary: Implement a new swapping API which does the following:

  1. Takes an exported program and torchrec serializer
  2. Constructs torchrec modules based on serialized metadata stored in the exported program
  3. Unlifts the exported program to a normal nn.Module
  4. Partitions out nodes and swaps them out with calls to the torchrec modules

Differential Revision: D63683421

angelayi avatar Oct 09 '24 18:10 angelayi