Skip to content
Snippets Groups Projects
Commit ee0f4a91 authored by Yaman Umuroglu's avatar Yaman Umuroglu
Browse files

[Transform] add infer_shapes

currently wrapper for ONNX shape inference, will change later
parent b2eec820
No related branches found
No related tags found
No related merge requests found
import onnx.shape_inference as si
def infer_shapes(model):
"""Ensure every tensor in the model has a specified shape (ValueInfo)."""
# currently just calls ONNX shape inference, but in the future we will
# have to handle shape inference for custom ops ourselves
model.model = si.infer_shapes(model.model)
# single-step operation, no need to call multiple times so return
# model_was_changed = false
return (model, False)
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment