Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add metadata for torch jit TracedModules. (#17640)
Summary: Pull Request resolved: #17640 Pull Request resolved: #17311 I've extended our model metadata framework in this diff to support traced modules as well. Re-used a lot of components from the previous implementation of ScriptModule metadata. Tracing is a little different from Scripting since you can't just create a subclass of TopLevelTraceModule (type returned by torch.jit.trace) and attach metadata the way we did for ScriptModule. As a result, I've introduced a separate API torch.fb.jit_trace which returns an instance of TracedModuleWithMetadata which is a subclass of TopLevelTracedModule. As a result, we can now attach metadata to this instance. Reviewed By: dzhulgakov Differential Revision: D14117966 fbshipit-source-id: 3eee5eef733cb8d6a219c02e2f41d08698eca326
- Loading branch information