Skip to content

Commit 76188a4

Browse files
srkreddy1238tqchen
authored andcommitted
[NNVM][TENSORFLOW] bugfix. (#2444)
1 parent a527b58 commit 76188a4

1 file changed

Lines changed: 3 additions & 0 deletions

File tree

nnvm/python/nnvm/frontend/tensorflow.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1193,6 +1193,9 @@ def from_tensorflow(self, graph, layout="NHWC", shape=None, outputs=None):
11931193
self._output_shapes[node.name] = \
11941194
[tensor_util.TensorShapeProtoToList( \
11951195
tensor_value.tensor_shape)]
1196+
elif shape and node.name in shape:
1197+
# Give priority to user argument.
1198+
self._output_shapes[node.name] = [shape[node.name]]
11961199
elif '_output_shapes' in attr:
11971200
self._output_shapes[node.name] = \
11981201
[tensor_util.TensorShapeProtoToList(tshape) \

0 commit comments

Comments
 (0)