Merge pull request #810 from pfnet/take-cheeze-patch-2

[onnx] Also check validation error in model load
pfnet · Mar 26, 2024 · 730085c · 730085c
2 parents 4979e30 + 08c20a5
commit 730085c
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/pytorch_pfn_extras/onnx/load.py b/pytorch_pfn_extras/onnx/load.py
@@ -24,8 +24,9 @@ def load_model(
     """
     try:
         return onnx.load_model(f, format=format, load_external_data=load_external_data)
-    except OSError as e:  # The ONNX may contain stripped large tensors.
+    except (OSError, onnx.checker.ValidationError) as e:  # The ONNX may contain stripped large tensors.
         if (load_external_data
+                and isinstance(e, OSError)
                 and json.loads(Path(e.filename).name)["type"] != "stripped"):
             raise
         warnings.warn(