Skip to content

T
tic
Commit 76188a43 by Siva Committed by Tianqi Chen
Options

[NNVM][TENSORFLOW] bugfix. (#2444)

parent a527b58e
Showing with 3 additions and 0 deletions
nnvm/python/nnvm/frontend/tensorflow.py
......@@ -1193,6 +1193,9 @@ class GraphProto(object):
self._output_shapes[node.name] = \
[tensor_util.TensorShapeProtoToList( \
tensor_value.tensor_shape)]
elif shape and node.name in shape:
# Give priority to user argument.
self._output_shapes[node.name] = [shape[node.name]]
elif '_output_shapes' in attr:
self._output_shapes[node.name] = \
[tensor_util.TensorShapeProtoToList(tshape) \
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment