ignore IValue for mypy check, since cannot import it from torch._C

pytorch · Jan 14, 2021 · 82d3a69 · 82d3a69
1 parent 070f4a9
commit 82d3a69
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/torch/onnx/utils.py b/torch/onnx/utils.py
@@ -17,8 +17,8 @@
 from torch._six import string_classes
 from torch.jit import _unique_state_dict
 from torch.onnx import ONNX_ARCHIVE_MODEL_PROTO_NAME, ExportTypes, OperatorExportTypes, TrainingMode
-from torch._C import ListType, OptionalType, _propagate_and_assign_input_shapes, _check_onnx_proto, IValue
-from typing import Union, Tuple, List, Dict
+from torch._C import ListType, OptionalType, _propagate_and_assign_input_shapes, _check_onnx_proto
+from typing import Union, Tuple, List
 
 
 # the flag to tell the user whether it's in the middle of ONNX export or not
@@ -29,7 +29,8 @@ def is_in_onnx_export():
     global __IN_ONNX_EXPORT
     return __IN_ONNX_EXPORT
 
-_params_dict: Dict[str, IValue] = {}
+# Skip check due to cannot import IValue from torch._C
+_params_dict = {}  # type: ignore
 
 @contextlib.contextmanager
 def select_model_mode_for_export(model, mode):