Skip to content

Commit 78d2d59

Browse files
authored
[95] Properly adds the shapes of the tensors to the graph. (#98)
* support for strings * removing type tags, they do not work with scala 2.10 * removing type tags, they do not work with scala 2.10 * removing type tags, they do not work with scala 2.10 * revert scala changes * revert scala changes * revert scala changes
1 parent b393bf3 commit 78d2d59

File tree

1 file changed

+3
-2
lines changed
  • src/main/python/tensorframes

1 file changed

+3
-2
lines changed

src/main/python/tensorframes/core.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,8 @@ def _add_graph(graph, builder, use_file=True):
4343
fname = d + "/proto.pb"
4444
builder.graphFromFile(fname)
4545
else:
46-
gser = graph.as_graph_def().SerializeToString()
46+
# Make sure that TF adds the shapes.
47+
gser = graph.as_graph_def(add_shapes=True).SerializeToString()
4748
gbytes = bytearray(gser)
4849
builder.graph(gbytes)
4950

@@ -55,7 +56,7 @@ def _add_shapes(graph, builder, fetches):
5556
# dimensions are unknown
5657
ph_names = []
5758
ph_shapes = []
58-
for n in graph.as_graph_def().node:
59+
for n in graph.as_graph_def(add_shapes=True).node:
5960
# Just the input nodes:
6061
if not n.input:
6162
op_name = n.name

0 commit comments

Comments
 (0)