Propagating onnx.export() parameters; requiring torch>=1.6

NVIDIA · Aug 31, 2020 · 7963b4c · 7963b4c
1 parent 6471d10
commit 7963b4c
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 13 deletions.
diff --git a/nemo/core/classes/exportable.py b/nemo/core/classes/exportable.py
@@ -47,13 +47,17 @@ class Exportable(ABC):
     """
 
     def export(
-        self,
-        output: str,
-        input_example=None,
-        output_example=None,
-        onnx_opset_version: int = 12,
-        try_script: bool = False,
-        set_eval: bool = True,
+            self,
+            output: str,
+            input_example=None,
+            output_example=None,
+            verbose=False,
+            export_params=True,
+            do_constant_folding=True,
+            keep_initializers_as_inputs=False,
+            onnx_opset_version: int = 12,
+            try_script: bool = False,
+            set_eval: bool = True,
     ):
         try:
             # Disable typechecks
@@ -134,10 +138,10 @@ def export(
                         output,
                         input_names=input_names,
                         output_names=output_names,
-                        verbose=False,
-                        export_params=True,
-                        do_constant_folding=True,
-                        keep_initializers_as_inputs=True,
+                        verbose=verbose,
+                        export_params=export_params,
+                        do_constant_folding=do_constant_folding,
+                        keep_initializers_as_inputs=keep_initializers_as_inputs,
                         dynamic_axes=dynamic_axes,
                         opset_version=onnx_opset_version,
                         example_outputs=_out_example,

diff --git a/requirements/requirements.txt b/requirements/requirements.txt
@@ -2,11 +2,11 @@ numpy>=1.18.2
 onnx>=1.7.0
 pytorch-lightning==0.9.0
 python-dateutil
-torch
+torch>=1.6.0
 wget
 wrapt
 ruamel.yaml
 scikit-learn
 omegaconf==2.0.1rc12
 hydra-core==1.0.0rc4
-transformers>=2.11.0
+transformers>=2.11.0