use TrackingAdapter for BoltNnExport

Differential Revision: D30615605 fbshipit-source-id: d4d4550b6d1da4c75945ba674fbdd49a842ec6a9

use TrackingAdapter for BoltNnExport
Differential Revision: D30615605 fbshipit-source-id: d4d4550b6d1da4c75945ba674fbdd49a842ec6a9
ecbe3e02 · Yanghan Wang · Facebook GitHub Bot · e62c0e4c · ecbe3e02
Commit ecbe3e02 authored Aug 31, 2021 by Yanghan Wang Committed by Facebook GitHub Bot Aug 31, 2021
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

d2go/export/torchscript.py d2go/export/torchscript.py +2 -2

No files found.
--- a/d2go/export/torchscript.py
+++ b/d2go/export/torchscript.py
@@ -142,8 +142,8 @@ def tracing_adapter_wrap_export(old_f):

 def tracing_adapter_wrap_load(old_f):
    def new_f(cls, save_path, **load_kwargs):
-        assert "inputs_schema" in load_kwargs
-        assert "outputs_schema" in load_kwargs
+        assert "inputs_schema" in load_kwargs, load_kwargs.keys()
+        assert "outputs_schema" in load_kwargs, load_kwargs.keys()
        inputs_schema = instantiate(load_kwargs.pop("inputs_schema"))
        outputs_schema = instantiate(load_kwargs.pop("outputs_schema"))
        traced_model = old_f(cls, save_path, **load_kwargs)