Skip to content

GitLab

Commit ecbe3e02 authored by Yanghan Wang's avatar Yanghan Wang Committed by Facebook GitHub Bot
Browse files

use TrackingAdapter for BoltNnExport 

Differential Revision: D30615605

fbshipit-source-id: d4d4550b6d1da4c75945ba674fbdd49a842ec6a9
parent e62c0e4c
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
d2go/export/torchscript.py
View file @ ecbe3e02
...@@ -142,8 +142,8 @@ def tracing_adapter_wrap_export(old_f): ...@@ -142,8 +142,8 @@ def tracing_adapter_wrap_export(old_f):
def tracing_adapter_wrap_load(old_f): def tracing_adapter_wrap_load(old_f):
def new_f(cls, save_path, **load_kwargs): def new_f(cls, save_path, **load_kwargs):
assert "inputs_schema" in load_kwargs assert "inputs_schema" in load_kwargs, load_kwargs.keys()
assert "outputs_schema" in load_kwargs assert "outputs_schema" in load_kwargs, load_kwargs.keys()
inputs_schema = instantiate(load_kwargs.pop("inputs_schema")) inputs_schema = instantiate(load_kwargs.pop("inputs_schema"))
outputs_schema = instantiate(load_kwargs.pop("outputs_schema")) outputs_schema = instantiate(load_kwargs.pop("outputs_schema"))
traced_model = old_f(cls, save_path, **load_kwargs) traced_model = old_f(cls, save_path, **load_kwargs)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment